Training a Vision model

In the ImageAnalyzer sample app, you saw that picking an image of a certain car would be classified as a sports car with a pretty low confidence score. You can train your own vision model that specializes in recognizing certain cars.

Collecting good training data for image classifiers is tough, because you have to make sure that you gather many pictures of your subjects from all sides and in many different environments. For instance, if all your car images feature cars that are next to trees, or on the road, the model might end up classifying anything with trees or a road to be a car. The only way to obtain a perfect training set is to experiment, tweak, and test.

Training a Vision model works slightly different from training a Natural Language model. You can't use a JSON file to feed your test data to the classifier. So, instead, you should create folders that contain your images where the folder name is the label you want to apply to each image inside that folder. The following screenshot is an example of a training set that contains two kinds of labels:

Once you have collected your set of training data, you can store it anywhere on your computer—for instance, on the desktop. You will then pass the path for your training data to your model training code as follows:

import CreateML
import Foundation

let dataUrl = URL(fileURLWithPath: "/path/to/trainingdata")
let source = MLImageClassifier.DataSource.labeledDirectories(at: dataUrl)
let classifier = try! MLImageClassifier(trainingData: source)
try! classifier.write(toFile: "~/Desktop/CarClassifier.mlmodel")

Again, you only need a couple of lines of code to train a model. That's how powerful CreateML is. If you want, you can quickly test your image classifier by dropping the .mlmodel file in the ImageAnalyzer project, and using that, instead of the MobileNet classifier that you used before.

Apart from the simple ways of training models, there are certain parameters that you can pass to the different CreateML classifiers. If you have trouble training your models properly, you could tweak some of the parameters that are used by CreateML. For instance, you could apply more iterations to your training set, so the model gains a deeper understanding of the training data.

As mentioned before in this chapter, machine learning is a subject that could span several books on its own, and even though CreateML makes training models straightforward and simple, it's not easy to train a robust model without any prior machine learning experience.