Revision history [back]

Here some information and pointers to train a neural network for object detection using TensorFlow on custom data.

Following are references I used to train with MobileNet SSD (using checkpoint weights) for one class on my data and locally (no cloud). I cannot guarantee that these instructions are optimal or the good way to do but it should work.

TensorFlow installation

See the official tutorial.

Installation should be quite easy, I did it on Ubuntu 16.04 using the virtualenv way. If you have a good nvidia GPU, it is worth the effort to install CUDA and cuDNN and install tensorflow-gpu.

What should be tricky is to install the latest display driver to be able to use the latest CUDA and cuDNN versions. On Ubuntu, you have to disable the X-server for instance and uninstall the drivers coming from the distribution I think. More information here: NVIDIA CUDA Installation Guide for Linux.

TensorFlow object detection API

• Documentation are here and here
• Installation instructions are here
• New API (July 13, 2018) doc is here

Guide to train on custom data

• This guide contains some useful information but uses now the legacy API (legacy since July 13, 2018)
• The procedure is always the same: labelize your data, you can use LabelImg for that
• Split your data into training and evaluation sets
• Labels are in PASCAL VOC format, see this answer to setup the directories
• Generate TFRecord data from XML annotation files
• Change the parameters (num_classes, batch_size, fine_tune_checkpoint, train_input_reader, eval_config, eval_input_reader) in ssd_mobilenet_v1_coco.config for your desired network
• Launch the training and the evaluation, see these instructions for the old API
• Launch Tensorboard: tensorboard --logdir=${PATH_TO_MODEL_DIRECTORY} and open in a browser the following address: http://localhost:6006
• You can then observe the loss curves and the results of the detection on your evaluation data

Export / Test

• Exporting a trained model for inference
• OpenCV TensorFlow Object Detection API compatibility

I have skipped some details to keep it short but all the information are available in the official documentation or in some tutorials. The procedure is always the same and once you know how to do it, it is quite straightforward.