TensorFlow Lite Model Maker

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 TasksTask Utility
Image Classification guideClassify images into predefined categories.
Text Classification guideClassify text into predefined categories.
Question Answer guideFind the answer in a certain context for a given question.

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.

# Load input data specific to an on-device ML app.
data = ImageClassifierDataLoader.from_folder('flower_photos/')
train_data, test_data = data.split(0.9)

# Customize the TensorFlow model.
model = image_classifier.create(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

Install a prebuilt pip package.

pip install tflite-model-maker