Just some months ago, Google released code for classifying images using neural networks. Some time later, they also released code to train your custom models, either from scratch or improving a baseline model. The baseline model in that case usually is a model trained on the ImageNet dataset.

To my understanding, there are currently two different ways to train a custom model:

  • The tensorflow/models repository contains several examples training models from scratch or improving models. The retraining in this case seems to include retraining on the full network, which is quite slow on my machine.
  • The tensorflow/tensorflow project contains an example file on retraining just the final layer of a model, but keeping all the rest from the baseline.

Due to the long runtime of the first approach on my computer (I did not want to apply for the cuDNN program), I decided for the second approach. The commentary section of that file already shows almost anything you have to do. Basically, you just have to compile it and then run it.

bazel build tensorflow/examples/image_retraining:retrain

The command line arguments are also given within the file as tf.app.flags.DEFINE_*. To pass them as argument, just preceed them with a double dash --.

bazel-bin/tensorflow/examples/image_retraining/retrain --image_dir ./images

This will then create a graph in your /tmp directory. Of course, you can also adjust the output path using the command line arguments --output_graph and --output_labels.

To classify another image with the new network you may use the label_image example. To compile it use:

bazel build tensorflow/examples/label_image:label_image

You can then classify a new image with the following command:

bazel-bin/tensorflow/examples/label_image/label_image \
--graph=/tmp/output_graph.pb --labels=/tmp/output_labels.txt \
--output_layer=final_result \
--image=./image.jpg

Of course, replace image.jpg with the real path to your image.