使用 TensorFlow Lite Model Maker 进行图像分类

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本 查看 TF Hub 模型

TensorFlow Lite Model Maker 库 简化了将 TensorFlow 神经网络模型调整和转换为特定输入数据的过程，以便在将此模型部署到设备上进行机器学习应用时使用。

此笔记本展示了一个端到端的示例，该示例使用 Model Maker 库来说明如何调整和转换常用的图像分类模型，以便在移动设备上对花卉进行分类。

先决条件

要运行此示例，我们首先需要安装几个必需的软件包，包括 GitHub 仓库 中的 Model Maker 软件包。

sudo apt -y install libportaudio2
pip install -q tflite-model-maker

导入所需的软件包。

import os

import numpy as np

import tensorflow as tf
assert tf.__version__.startswith('2')

from tflite_model_maker import model_spec
from tflite_model_maker import image_classifier
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.config import QuantizationConfig
from tflite_model_maker.image_classifier import DataLoader

import matplotlib.pyplot as plt

简单的端到端示例

获取数据路径

让我们获取一些图像来试用这个简单的端到端示例。数百张图像对于 Model Maker 来说是一个良好的开端，而更多的数据可以获得更高的准确率。

您可以用自己的图像文件夹替换 image_path。至于将数据上传到 Colab，您可以在左侧边栏中找到上传按钮，如下图中红色矩形所示。只需尝试上传一个 zip 文件并解压缩它即可。根文件路径是当前路径。

Upload File

如果您不想将图像上传到云端，可以尝试按照 GitHub 上的 指南 在本地运行库。

运行示例

此示例仅包含 4 行代码，如下所示，每行代码代表整个过程中的一个步骤。

步骤 1. 加载特定于设备上机器学习应用的输入数据。将其拆分为训练数据和测试数据。

data = DataLoader.from_folder(image_path)
train_data, test_data = data.split(0.9)

步骤 2. 自定义 TensorFlow 模型。

model = image_classifier.create(train_data)

步骤 3. 评估模型。

loss, accuracy = model.evaluate(test_data)

步骤 4. 导出到 TensorFlow Lite 模型。

在这里，我们使用 元数据 导出 TensorFlow Lite 模型，该元数据提供模型描述的标准。标签文件嵌入在元数据中。图像分类任务的默认训练后量化技术是全整数量化。

您可以在左侧边栏中下载它，就像上传部分一样，供您自己使用。

model.export(export_dir='.')

完成这 4 个简单的步骤后，我们可以在设备上应用程序中进一步使用 TensorFlow Lite 模型文件，例如 图像分类 参考应用程序。

详细过程

目前，我们支持几种模型，例如 EfficientNet-Lite* 模型、MobileNetV2、ResNet50 作为图像分类的预训练模型。但只需几行代码，就可以非常灵活地将新的预训练模型添加到此库中。

以下将逐步介绍此端到端示例，以显示更多细节。

步骤 1：加载特定于设备上机器学习应用的输入数据

花卉数据集包含 3670 张图像，属于 5 个类别。下载数据集的存档版本并解压缩它。

数据集具有以下目录结构

flower_photos
|__ daisy
    |______ 100080576_f52e8ee070_n.jpg
    |______ 14167534527_781ceb1b7a_n.jpg
    |______ ...
|__ dandelion
    |______ 10043234166_e6dd915111_n.jpg
    |______ 1426682852_e62169221f_m.jpg
    |______ ...
|__ roses
    |______ 102501987_3cdb8e5394_n.jpg
    |______ 14982802401_a3dfb22afb.jpg
    |______ ...
|__ sunflowers
    |______ 12471791574_bb1be83df4.jpg
    |______ 15122112402_cafa41934f.jpg
    |______ ...
|__ tulips
    |______ 13976522214_ccec508fe7.jpg
    |______ 14487943607_651e8062a1_m.jpg
    |______ ...

image_path = tf.keras.utils.get_file(
      'flower_photos.tgz',
      'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
      extract=True)
image_path = os.path.join(os.path.dirname(image_path), 'flower_photos')

使用 DataLoader 类加载数据。

至于 from_folder() 方法，它可以从文件夹加载数据。它假设同一类别的图像位于同一个子目录中，子文件夹名称是类别名称。目前，支持 JPEG 编码的图像和 PNG 编码的图像。

data = DataLoader.from_folder(image_path)

将其拆分为训练数据 (80%)、验证数据 (10%，可选) 和测试数据 (10%)。

train_data, rest_data = data.split(0.8)
validation_data, test_data = rest_data.split(0.5)

显示 25 个带有标签的图像示例。

plt.figure(figsize=(10,10))
for i, (image, label) in enumerate(data.gen_dataset().unbatch().take(25)):
  plt.subplot(5,5,i+1)
  plt.xticks([])
  plt.yticks([])
  plt.grid(False)
  plt.imshow(image.numpy(), cmap=plt.cm.gray)
  plt.xlabel(data.index_to_label[label.numpy()])
plt.show()

步骤 2：自定义 TensorFlow 模型

基于加载的数据创建自定义图像分类模型。默认模型为 EfficientNet-Lite0。

model = image_classifier.create(train_data, validation_data=validation_data)

查看详细的模型结构。

model.summary()

步骤 3：评估自定义模型

评估模型的结果，获取模型的损失和准确率。

loss, accuracy = model.evaluate(test_data)

我们可以绘制 100 张测试图像的预测结果。红色标记的预测标签表示预测错误的结果，其他颜色表示预测正确的结果。

# A helper function that returns 'red'/'black' depending on if its two input
# parameter matches or not.
def get_label_color(val1, val2):
  if val1 == val2:
    return 'black'
  else:
    return 'red'

# Then plot 100 test images and their predicted labels.
# If a prediction result is different from the label provided label in "test"
# dataset, we will highlight it in red color.
plt.figure(figsize=(20, 20))
predicts = model.predict_top_k(test_data)
for i, (image, label) in enumerate(test_data.gen_dataset().unbatch().take(100)):
  ax = plt.subplot(10, 10, i+1)
  plt.xticks([])
  plt.yticks([])
  plt.grid(False)
  plt.imshow(image.numpy(), cmap=plt.cm.gray)

  predict_label = predicts[i][0][0]
  color = get_label_color(predict_label,
                          test_data.index_to_label[label.numpy()])
  ax.xaxis.label.set_color(color)
  plt.xlabel('Predicted: %s' % predict_label)
plt.show()

如果准确率不满足应用程序要求，可以参考 高级用法 探索其他方案，例如切换到更大的模型、调整重新训练参数等。

步骤 4：导出到 TensorFlow Lite 模型

使用 元数据 将训练后的模型转换为 TensorFlow Lite 模型格式，以便您可以在设备上的机器学习应用程序中使用它。标签文件和词汇文件嵌入在元数据中。默认的 TFLite 文件名为 model.tflite。

在许多设备上的机器学习应用程序中，模型大小是一个重要因素。因此，建议您对模型进行量化，使其更小，并可能运行得更快。图像分类任务的默认后训练量化技术是全整数量化。

model.export(export_dir='.')

有关如何将 TensorFlow Lite 模型集成到移动应用程序的更多详细信息，请参阅图像分类 示例指南。

可以使用 ImageClassifier API（属于 TensorFlow Lite 任务库）将此模型集成到 Android 或 iOS 应用程序中。

允许的导出格式可以是以下格式之一或列表：

默认情况下，它只导出带有元数据的 TensorFlow Lite 模型。您也可以选择性地导出不同的文件。例如，仅导出标签文件，如下所示：

model.export(export_dir='.', export_format=ExportFormat.LABEL)

您还可以使用 evaluate_tflite 方法评估 tflite 模型。

model.evaluate_tflite('model.tflite', test_data)

高级用法

create 函数是此库的关键部分。它使用与 教程 相似的预训练模型进行迁移学习。

create 函数包含以下步骤：

  1. 根据参数 validation_ratio 和 test_ratio 将数据分成训练、验证、测试数据。validation_ratio 和 test_ratio 的默认值为 0.1 和 0.1。
  2. 从 TensorFlow Hub 下载 图像特征向量 作为基础模型。默认的预训练模型为 EfficientNet-Lite0。
  3. 添加一个分类器头部，并在头部层和预训练模型之间添加一个具有 dropout_rate 的 Dropout 层。默认的 dropout_rate 是 TensorFlow Hub 的 make_image_classifier_lib 中的默认 dropout_rate 值。
  4. 预处理原始输入数据。目前，预处理步骤包括将每个图像像素的值归一化为模型输入比例，并将其调整为模型输入大小。EfficientNet-Lite0 的输入比例为 [0, 1]，输入图像大小为 [224, 224, 3]。
  5. 将数据馈送到分类器模型。默认情况下，训练参数（例如训练轮次、批次大小、学习率、动量）是 TensorFlow Hub 的 make_image_classifier_lib 中的默认值。仅训练分类器头部。

在本节中，我们将介绍几个高级主题，包括切换到不同的图像分类模型、更改训练超参数等。

自定义 TensorFlow Lite 模型上的后训练量化

后训练量化 是一种转换技术，可以减小模型大小和推理延迟，同时提高 CPU 和硬件加速器的推理速度，而模型精度略有下降。因此，它被广泛用于优化模型。

Model Maker 库在导出模型时应用默认的后训练量化技术。如果您想自定义后训练量化，Model Maker 也支持使用 QuantizationConfig 的多种后训练量化选项。以 float16 量化为例。首先，定义量化配置。

config = QuantizationConfig.for_float16()

然后，我们使用此配置导出 TensorFlow Lite 模型。

model.export(export_dir='.', tflite_filename='model_fp16.tflite', quantization_config=config)

在 Colab 中，您可以从左侧边栏下载名为 model_fp16.tflite 的模型，与上面提到的上传部分相同。

更改模型

更改为此库支持的模型。

此库目前支持 EfficientNet-Lite 模型、MobileNetV2、ResNet50。 EfficientNet-Lite 是一个图像分类模型系列，可以实现最先进的精度，并且适合边缘设备。默认模型为 EfficientNet-Lite0。

我们可以通过将参数 model_spec 设置为 create 方法中的 MobileNetV2 模型规范来切换到 MobileNetV2 模型。

model = image_classifier.create(train_data, model_spec=model_spec.get('mobilenet_v2'), validation_data=validation_data)

评估新重新训练的 MobileNetV2 模型，以查看测试数据的准确率和损失。

loss, accuracy = model.evaluate(test_data)

更改 TensorFlow Hub 中的模型

此外，我们还可以切换到其他新的模型，这些模型输入图像并输出具有 TensorFlow Hub 格式的特征向量。

以 Inception V3 模型为例，我们可以定义 inception_v3_spec，它是一个 image_classifier.ModelSpec 对象，包含 Inception V3 模型的规范。

我们需要指定模型名称 name 和 TensorFlow Hub 模型的 URL uri。同时，input_image_shape 的默认值为 [224, 224]。对于 Inception V3 模型，我们需要将其更改为 [299, 299]。

inception_v3_spec = image_classifier.ModelSpec(
    uri='https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1')
inception_v3_spec.input_image_shape = [299, 299]

然后，通过将参数 model_spec 设置为 create 方法中的 inception_v3_spec，我们可以重新训练 Inception V3 模型。

其余步骤完全相同，最终我们可以获得一个自定义的 InceptionV3 TensorFlow Lite 模型。

更改您自己的自定义模型

如果我们想使用 TensorFlow Hub 中没有的自定义模型，我们应该在 TensorFlow Hub 中创建和导出 ModelSpec。

然后开始像上面那样定义 ModelSpec 对象。

更改训练超参数

我们还可以更改训练超参数，例如 epochs、dropout_rate 和 batch_size，这些超参数会影响模型的准确率。您可以调整的模型参数有：

  • epochs：更多的轮次可以实现更好的精度，直到收敛，但训练过多的轮次可能会导致过拟合。
  • dropout_rate：Dropout 的比率，避免过拟合。默认情况下为 None。
  • batch_size：一次训练步骤中使用的样本数量。默认情况下为 None。
  • validation_data：验证数据。如果为 None，则跳过验证过程。默认情况下为 None。
  • train_whole_model：如果为 True，则 Hub 模块将与顶部的分类层一起训练。否则，只训练顶部的分类层。默认情况下为 None。
  • learning_rate：基本学习率。默认情况下为 None。
  • momentum：传递给优化器的 Python 浮点数。仅在 use_hub_library 为 True 时使用。默认情况下为 None。
  • shuffle：布尔值，指示是否应对数据进行洗牌。默认情况下为 False。
  • use_augmentation：布尔值，使用数据增强进行预处理。默认情况下为 False。
  • use_hub_library：布尔值，使用来自 tensorflow hub 的 make_image_classifier_lib 重新训练模型。此训练管道可以为具有许多类别的复杂数据集实现更好的性能。默认情况下为 True。
  • warmup_steps：学习率预热计划的预热步骤数。如果为 None，则使用默认的预热步骤，即两个轮次中的总训练步骤。仅在 use_hub_library 为 False 时使用。默认情况下为 None。
  • model_dir：可选，模型检查点文件的存储位置。仅在 use_hub_library 为 False 时使用。默认情况下为 None。

默认情况下为 None 的参数（例如 epochs）将在 TensorFlow Hub 库的 make_image_classifier_lib 或 train_image_classifier_lib 中获取具体的默认参数。

例如，我们可以使用更多轮次进行训练。

model = image_classifier.create(train_data, validation_data=validation_data, epochs=10)

评估使用 10 个训练轮次重新训练的模型。

loss, accuracy = model.evaluate(test_data)

阅读更多

您可以阅读我们的 图像分类 示例以了解技术细节。有关更多信息，请参考：