TensorFlow Lite 模型制作工具

概览

TensorFlow Lite Model Maker 库可简化使用自定义数据集训练 TensorFlow Lite 模型的过程。它使用迁移学习来减少所需的训练数据量并缩短训练时间。

支持的任务

Model Maker 库目前支持以下机器学习任务。点击以下链接,查看有关如何训练模型的指南。

支持的任务 任务实用程序
图片分类:教程API 将图片归入预定义的类别。
对象检测:教程API 实时检测对象。
文本分类:教程API 将文本归类到预定义的类别中。
BERT 问题解答:教程API 使用 BERT 在特定上下文中查找给定问题的答案。
音频分类:教程API 将音频归类到预定义的类别中。
建议:演示API 根据设备端场景的上下文信息推荐内容。
搜索者:教程API 在数据库中搜索相似的文本或图片。

如果您的任务不受支持,请先使用 TensorFlow 通过迁移学习(按照图片文本音频等指南)重新训练 TensorFlow 模型,或从头开始训练模型,然后将其转换为 TensorFlow Lite 模型。

端到端示例

借助 Model Maker,您只需几行代码即可使用自定义数据集训练 TensorFlow Lite 模型。例如,以下是训练图片分类模型的步骤。

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/')

如需了解详情,请参阅图片分类指南

安装

您可以通过两种方式安装 Model Maker。

  • 安装预构建的 pip 软件包。
pip install tflite-model-maker

如果您想安装每夜版,请按照以下命令操作:

pip install tflite-model-maker-nightly
  • 从 GitHub 克隆源代码并进行安装。
git clone https://github.com/tensorflow/examples
cd examples/tensorflow_examples/lite/model_maker/pip_package
pip install -e .

TensorFlow Lite Model Maker 依赖于 TensorFlow pip 软件包。对于 GPU 驱动程序,请参阅 TensorFlow 的 GPU 指南安装指南

Python API 参考文档

您可以在 API 参考文档中找到 Model Maker 的公开 API。