Wolfram Language

Train a Custom Image Classifier

Some classification tasks that are difficult for humans turn out to be quite easy for neural nets. For instance, many people cannot remember the difference between a camel and a dromedary. Using a simple fine-tuning procedure, a neural net can help.

On the Wolfram Neural Net Repository, there are several models trained to identify the main object in an image. Here is one model trained on the ImageNet classification dataset.

In its original state, it is not good at differentiating a camel from a dromedary: only one animal is present in the ImageNet training data.

However, the model can be retrained to perform better on this specific task. To do the retraining, remove the final classification layers and replace them with a two-class classifier.

New training data can be gathered using WebImageSearch.

Using NetTrain, the net is trained on the new data, saving 10% for validation measurements. Only the final classification layer is trained.

The new model gives a much better classification.

Using the same procedure, retrain three other classification architectures trained on ImageNet and compare the accuracy of classification.

For a simple task like this, the accuracy is similar and there is no need for the more powerful and slower models. The small ResNet is a good tradeoff between accuracy and speed.

Related Examples

ja