Skip to content

Transfer learning with SSD MobileNet v1

David Salek edited this page Jan 16, 2018 · 1 revision

A more ambitious goal is to use transfer learning to turn a pre-trained SSD MobileNet v1 model into a Santa detector.

The inspiration for doing this comes from the excellent blog post by Dat Tran below. In this post, Dat follows the instructions from the TensorFlow Object Detection API documentation to recognise custom objects. https://towardsdatascience.com/how-to-train-your-own-object-detector-with-tensorflows-object-detector-api-bec72ecfe1d9

Although I managed to train a new model using transfer learning, it does not perform well due to an insufficient dataset. Gathering and annotating images is a time-consuming task and in the end I only decided to test the instructions on a handful of images.


Interactive Docker container

Run an interactive shell using the ubuntu:16.04 Docker image.

docker run -i -t -v $(pwd):/santa ubuntu:16.04 /bin/bash

Install TensorFlow Object Detection API and download a pre-trained SSD MobileNet v1** model.

apt-get update
apt-get install -y git
git clone https://github.com/salekd/rpizero_smart_camera3.git
git clone https://github.com/tensorflow/models.git

apt-get install -y wget
wget http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_11_06_2017.tar.gz
tar -zxvf ssd_mobilenet_v1_coco_11_06_2017.tar.gz

apt-get install -y python-pip python-dev build-essential
pip install tensorflow
pip install Pillow
pip install matplotlib

apt-get install -y protobuf-compiler
cd /models/research/
protoc object_detection/protos/*.proto --python_out=.
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim

apt-get install -y python-tk

Prepare a dataset

The following links describe how to prepare a dataset for training.

I created this script that uses a csv file to read the training images and the bounding box coordinates for Santa: https://github.com/salekd/santa/blob/master/transfer_learning/create_santa_tf_record.py

python create_santa_tf_record.py

Training

How to train a model, starting from a pre-trained checkpoint, is described here: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/configuring_jobs.md

Since I am using SSD MobileNet v1 for the Santa recognition, I start from the original configuration file: https://github.com/tensorflow/models/blob/master/research/object_detection/samples/configs/ssd_mobilenet_v1_coco.config

The modified version of the configuration file and a file defining the new label for Santa are part of this repository.

Run training as described here: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/running_locally.md

python /models/research/object_detection/train.py \
--logtostderr \
--pipeline_config_path=ssd_mobilenet_v1_coco.config\
--train_dir=train_dir

Export the model as described here: https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/exporting_models.md

python /models/research/object_detection/export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path train_dir/pipeline.config \
--trained_checkpoint_prefix train_dir/model.ckpt-1000 \
--output_directory ssd_mobilenet_santa

You can use the new model for inference, take a look at the example here: https://github.com/salekd/rpizero_smart_camera3/blob/master/object_detection_test.py

You can’t perform that action at this time.