How to retrain an object detection model with a custom training set

This tutorial will show how to retrain a mobilenetV2 model with a custom set of images in 15 minutes.

Published in

Analytics Vidhya

6 min readJan 13, 2021

To train an object detection model from scratch will require long hours of model training. To save time, the simplest approach would be to use an already trained model and retrain it to detect your custom objects. This process is called “transfer learning”.

Why mobilenetV2?

Object detection model performance continues to improve. As of 2020, MobilenetV2 is the fastest object detection model which can run in a web browser.

There are other models for object detection such as YOLO and more. You can find a detailed object detection model review here.

About Google Colab

MobilenetV2 is a model based on TensorFlow, therefore we will execute commands in the google colab environment with an open-source object_detection API based on TensorFlow. Google colab is free, but it deletes your output files after a short while, which is an unfortunate downside.

You can create your own Linux server/docker running object_detection API with all libraries already installed. For details see this blog.

Agenda for this tutorial

In this tutorial we are going to retrain mobilenetV2 model to detect a Pickachu toy. We will execute these 10 steps:

Step 1: Clone object detection API libraries in your colab notebook
Step 2 to 4 : Install object detection API and import its functions
Step 5: Import your training set of images to colab
Step 6: Set training configuration
Step 7: Load pre-trained mobilenetv2 model with its weights and latest checkpoint
Step 8: Train our new model with our custom training set
Step 9: Export our trained model to “SavedModel” format
Step 10: Use our model to detect Pickachu in test image.

Here you can find the full notebook, just run it in colab (It takes about 15 minutes to run all the steps).

File system map in colab environment

The tricky part working with the object detection API library is to keep track of all input and output files.

To make life a bit easier, here is a file map of colab environment with all the files we are going to use in this tutorial.

The folders are numbered according to execution order and icons indicate whether the directory is part of input or output of the training process.
Files will be loaded accordingly to execution steps of this tutorial, therefore not all files will be there from the start.

Execution Steps

Run the following steps in colab notebook (full notebook is available here).

Step 1: Clone TensorFlow2 object detection API folder to your colab:

Step 2: Install object detection API in colab:

Note: unfortunately, for some reason, you need to run the step twice to fix errors in installation.

Step 3: Import the TensorFlow libraries to colab:

Step 4: Run this command to test if object detection API is installed properly:

!python /content/models/research/object_detection/builders/model_builder_tf2_test.py

Step 5: Prepare custom set of images for training. I’m using Roboflow in this example to create the TFRecord file from the training images.

You can also create a labeled training data set with available open-source libraries check this article.
Roboflow gives you limited free access to try its automated process, which is straightforward and simple.
An example of TFRecord files is available at my git repository.
In this tutorial I used only 10 images to train the model. Adding images to training set, in most cases, will improve the detection.

Downloaded files will go to /content/train/ folder in colab environment.

Step 6: Set configuration parameters for training the model:

Step 7: Load trained model weights from TensorFlow site, and the custom configuration file which I’ve prepared for this example:

https://raw.githubusercontent.com/KostaMalsev/RetrainModelExample/main/ssd_mobilenet_v2_320x320_coco17_tpu-8.config

Note: Make sure you use short names for filenames and folders and paths. Training process fails if you use long path names in your configuration file!
Here are some notes regarding pipeline_file configuration file (ssd_mobilenet_v2_320x320_coco17_tpu-8.config):

#Choose type ssd_mobilenet_v2_keras
feature_extractor {                             
  type: "ssd_mobilenet_v2_keras"
..#Training batch size: if you have more memory on server,by increasing batch_size will run the training faster.
batch_size: 16#The path to pretrained model checkpoint:(choose short names) fine_tune_checkpoint: "/content/deploy/mobilnetv2/checkpoint/ckpt-0"#Number of training steps:
num_steps: 1800#Type of checkpoint when loading from pretrained model: fine_tune_checkpoint_type: "detection"#Defines the location of train set of images TFRecords and label maps
train_input_reader {                                        label_map_path: "/content/train/Toy_label_map.pbtxt"                         tf_record_input_reader {                                    input_path: "/content/train/Toy.tfrecord"                         }
..

Step 8: Run the training process:

The following parameters , together with pipeline_file configuration will be used in the training process.

The configuration values in pipline_file (ssd_mobilenet_v2_320x320_coco17_tpu-8.config) will overwrite the inline parameters.

pipeline_file: defined above in writing custom training configuration. (ssd_mobilenet_v2_320x320_coco17_tpu-8.config)model_dir: the location where logs and saved model checkpoints will save to.num_train_steps: how long to train for.num_eval_steps: perform eval on validation set after this many steps

Note: The training process takes several minutes, depends on your machine.
In most cases the longer the training the better are the detection results.
In this example I got a sound result with 1100 steps with loss=0.159. It took ~10 minutes of running in colab.

Look for checkpoint (ckpt ) files in /content/training/train folder.

You can stop the training process if the loss function result satisfies your requirements. The loss function result can be found in the notebook log, look for the following line:

Step 1100 per-step time 0.478s loss=0.159

From my experience, loss function result with a value lower than0.2 is enough for a demo application.

Step 9: Save the retrained model:

Exported model will be stored at output_directory. In this example the model will be exported to /content/fine_tuned_model.

That’s it.

Now you have a retrained model .pb file with checkpoint file which you can use later to for object detection.

Testing our retrained model:

To test our retrained model execute 4 additional steps:

Step A: Import the images you want to test to the /content/test folder. I used a Pickachu image from Pinterest.

#Import your test images to colab. 
#I use Pinterest to store the the images.%mkdir /content/test/%cd /content/test/!curl -L "https://i.pinimg.com/564x/f5/46/c4/f546c47505e1f5f8d17f8458d641b262.jpg" > test.jpeg;

Step B: Import visualization libraries and object detection functions:

Step C: Import our retrained model, its last training checkpoint and map object labels. For this example I chose ckpt-2 (second checkpoint).

Step D: Perform object detection on test images:

And here is the result:)

You can find the full notebook here.

Summary

In this tutorial we retrained an object detection model “mobilenetV2” with object detection API library and training images of our choosing.
The retrained model was exported to frozen format .pb file also known as “SavedModel” format.
We tested our new retrained object detection model on test images. Training and test images were converted to TFRecord format using Roboflow platform.

Next steps

We can use the retrained mobilenetV2 to perform object detection on live streaming camera in web browser, for tutorial see my article here.

Thank you for reading!