In this post, we walk through the steps to train and export a custom TensorFlow Lite object detection model with your own dataset to detect your own custom objects.

If you need a fast model on lower-end hardware, this post is for you. Whether for mobile phones or IoT devices, optimization is an especially important last step before deployment due to their lower performance.

In this tutorial, we will train an object detection model on custom data and convert it to TensorFlow Lite for deployment. We’ll conclude with a .tflite file that you can use in the official TensorFlow Lite Android Demo, iOS Demo, or Raspberry Pi Demo.

A Note about Custom Data

If you don’t have a dataset, you can follow along with a free Public Blood Cell Detection Dataset. If you have an unlabeled dataset, you can learn how to label with best practices with this post on how to label images with CVAT.

Let's dive in!

Open up this Colab Notebook to Train TensorFlow Lite Model.

What is TensorFlow Lite?

TensorFlow Lite is the official TensorFlow framework for on-device inference, meant to be used for small devices to avoid a round-trip to the server. This has many advantages, such as greater capacity for real-time detection, increased privacy, and not requiring an internet connection.

Note TensorFlow Lite isn’t for training models. It’s for bringing them to production. Luckily, the associated Colab Notebook for this post contains all the code to both train your model in TensorFlow and bring it to production in TensorFlow Lite.

Preparing Object Detection Data

TensorFlow models need data in the TFRecord format to train. Luckily, Roboflow converts any dataset into this format for us.

Note: if you have unlabeled data, you will first need to draw bounding boxes around your object in order to teach the computer to detect them. See this guide on how to label images in CVAT.

Once you have labeled data, if you haven’t yet, sign up for a free account with Roboflow. Then upload your dataset.

Drag and drop to upload data to Roboflow in any format

After upload you will be prompted to choose options including preprocessing and augmentations.

The augmentation and preprocessing selections for our BCCD dataset

After selecting these options, click Generate and then Download. You will be prompted to choose a data format for your export. Choose Tensorflow TFRecord.

Note: in addition to converting data to TFRecords in Roboflow, you can check the health of your dataset, make preprocessing decisions, and generate more data with augmentations so you can avoid overfitting and spend less time labeling more data.

Selecting TensorFlow TFRecord export format from Roboflow

After export, you will receive a curl link to download your data into our training notebook. There will be two links in the notebook to replace. The first is the TFRecord link you just generated. For the second link, repeat the steps for COCO JSON (in the dropdown as JSON > COCO).

Downloading our training data as TensorFlow TFRecord in the TFLite Colab Notebook

Lastly, we map our training data files to variables for use in our training pipeline configuration.

Training Custom TensorFlow Model

Because TensorFlow Lite lacks training capabilities, we will be training a TensorFlow 1 model beforehand: MobileNet Single Shot Detector (v2). Instead of writing the training from scratch, the training in this tutorial is based on a previous post: How to Train a TensorFlow MobileNet Object Detection Model. It explains the reason for choosing the MobileNet architecture:

“This specific architecture, researched by Google, is optimized for lightweight inference, enabling it to perform well natively on compute-constrained mobile and embedded devices”

We will not dive deeply into details here. To learn more about training, refer to the tutorial to train MobileNet and provided MobileNet Object Detection Colab Notebook.

From a high level, in order to train our custom object detection model, we take the following steps in the Colab Notebook to Train TensorFlow Lite Model:

  • Install TensorFlow object detection library and dependencies
  • Import dataset from Roboflow in TFRecord format
  • Write custom model configuration
  • Start custom TensorFlow object detection training job
Training our custom TensorFlow Lite object detection model ⏰
  • Export frozen inference graph in .pb format
  • Make inferences on test images to make sure our detector is functioning
Our custom cell detector inferring on a test image

Converting a SavedModel to TensorFlow Lite

Once you have a SavedModel (from the provided Colab Notebook or your own source), use the TensorFlow Lite Converter. Google recommends using its Python API, though it provides a Command Line Tool to do the trick as well.

We use the command line converter in the notebook because it’s simpler.

Converting our .pb model to .tflite with the command line converter

Be sure to set the input shape as desired for deployment. Smaller input shapes will run faster, but will be less performant.

Deployment Your Custom TensorFlow Lite Model

Lastly, we download our TensorFlow Lite model out of the Colab Notebook.

To deploy your model on device, check out the official TensorFlow Lite Android Demo, iOS Demo, or Raspberry Pi Demo. You can simply clone one of these repositories, drop in your .tflite file, and build according to the repo’s README.

You’re ready to run on-device inference 💥

As always - happy detecting!