Overview
The TensorFlow Lite Model Maker library simplifies the process of training a TensorFlow Lite model using custom dataset. It uses transfer learning to reduce the amount of training data required and shorten the training time.
Supported Tasks
The Model Maker library currently supports the following ML tasks. Click the links below for guides on how to train the model.
| Supported Tasks | Task Utility |
|---|---|
| Image Classification: tutorial, api | Classify images into predefined categories. |
| Object Detection: tutorial, api | Detect objects in real time. |
| Text Classification: tutorial, api | Classify text into predefined categories. |
| BERT Question Answer: tutorial, api | Find the answer in a certain context for a given question with BERT. |
| Audio Classification: tutorial, api | Classify audio into predefined categories. |
| Recommendation: demo, api | Recommend items based on the context information for on-device scenario. |
| Searcher: tutorial, api | Search for similar text or image in a database. |
If your tasks are not supported, please first use TensorFlow to retrain a TensorFlow model with transfer learning (following guides like images, text, audio) or train it from scratch, and then convert it to TensorFlow Lite model.
End-to-End Example
Model Maker allows you to train a TensorFlow Lite model using custom datasets in just a few lines of code. For example, here are the steps to train an image classification model.
from tflite_model_maker import image_classifier
from tflite_model_maker.image_classifier import DataLoader
# Load input data specific to an on-device ML app.
data = DataLoader.from_folder('flower_photos/')
train_data, test_data = data.split(0.9)
# Customize the TensorFlow model.
model = image_classifier.create(train_data)
# Evaluate the model.
loss, accuracy = model.evaluate(test_data)
# Export to Tensorflow Lite model and label file in `export_dir`.
model.export(export_dir='/tmp/')
For more details, see the image classification guide.
Installation
There are two ways to install Model Maker.
- Install a prebuilt pip package.
pip install tflite-model-maker
If you want to install nightly version, please follow the command:
pip install tflite-model-maker-nightly
- Clone the source code from GitHub and install.
git clone https://github.com/tensorflow/examples
cd examples/tensorflow_examples/lite/model_maker/pip_package
pip install -e .
TensorFlow Lite Model Maker depends on TensorFlow pip package. For GPU drivers, please refer to TensorFlow's GPU guide or installation guide.
Python API Reference
You can find out Model Maker's public APIs in API reference.