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.
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 guide | Classify images into predefined categories. |
Text Classification guide | Classify text into predefined categories. |
Question Answer guide | Find the answer in a certain context for a given question. |
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.
Install a prebuilt pip package.
pip install tflite-model-maker