在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

在本 Colab 笔记本中，您将学习如何使用 TensorFlow Lite Model Maker 库训练一个自定义目标检测模型，该模型能够在移动设备上检测图像中的沙拉。

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

您将使用公开可用的Salads 数据集，该数据集是从 Open Images Dataset V4 创建的。

数据集中的每个图像都包含标记为以下类别之一的对象

• 烘焙食品
• 奶酪
• 沙拉
• 海鲜
• 西红柿

数据集包含指定每个对象位置的边界框，以及对象的标签。

以下是从数据集中提取的示例图像

先决条件

安装所需的软件包

首先安装所需的软件包，包括来自 GitHub 仓库 的 Model Maker 软件包，以及您将用于评估的 pycocotools 库。

sudo apt -y install libportaudio2
pip install -q --use-deprecated=legacy-resolver tflite-model-maker
pip install -q pycocotools
pip install -q opencv-python-headless==4.1.2.30
pip uninstall -y tensorflow && pip install -q tensorflow==2.8.0

导入所需的软件包。

import numpy as np
import os

from tflite_model_maker.config import QuantizationConfig
from tflite_model_maker.config import ExportFormat
from tflite_model_maker import model_spec
from tflite_model_maker import object_detector

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

tf.get_logger().setLevel('ERROR')
from absl import logging
logging.set_verbosity(logging.ERROR)

准备数据集

这里您将使用与 AutoML 快速入门 相同的数据集。

Salads 数据集位于：gs://cloud-ml-data/img/openimage/csv/salads_ml_use.csv。

该数据集包含 175 张用于训练的图像、25 张用于验证的图像和 25 张用于测试的图像。数据集包含五个类别：Salad、Seafood、Tomato、Baked goods、Cheese。

数据集以 CSV 格式提供。

TRAINING,gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg,Salad,0.0,0.0954,,,0.977,0.957,,
VALIDATION,gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg,Seafood,0.0154,0.1538,,,1.0,0.802,,
TEST,gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg,Tomato,0.0,0.655,,,0.231,0.839,,

• 每行对应于一个位于较大图像内的对象，每个对象被专门指定为测试、训练或验证数据。您将在本笔记本的后续阶段了解有关其含义的更多信息。
• 这里包含的三行表示位于同一图像内的三个不同对象，该图像位于 gs://cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg。
• 每行都有不同的标签：Salad、Seafood、Tomato 等。
• 使用左上角和右下角顶点为每个图像指定边界框。

以下是这三行的可视化。

如果您想了解有关如何准备自己的 CSV 文件以及创建有效数据集的最低要求的更多信息，请参阅 准备您的训练数据 指南以获取更多详细信息。

如果您不熟悉 Google Cloud，您可能想知道 gs:// URL 的含义。它们是存储在 Google Cloud Storage (GCS) 上的文件的 URL。如果您将 GCS 上的文件设为公开或 对您的客户端进行身份验证，Model Maker 可以像读取本地文件一样读取这些文件。

但是，您无需将图像保存在 Google Cloud 上即可使用 Model Maker。您可以在 CSV 文件中使用本地路径，Model Maker 将正常工作。

快速入门

训练对象检测模型需要六个步骤。

步骤 1. 选择对象检测模型架构。

本教程使用 EfficientDet-Lite0 模型。EfficientDet-Lite[0-4] 是从 EfficientDet 架构派生的一系列移动/物联网友好型对象检测模型。

以下是每个 EfficientDet-Lite 模型与其他模型相比的性能。

* 整数量化模型的大小。
** 在 Pixel 4 上使用 CPU 上的 4 个线程测量的延迟。
*** 平均精度是 COCO 2017 验证数据集上的 mAP（平均精度均值）。

spec = model_spec.get('efficientdet_lite0')

步骤 2. 加载数据集。

Model Maker 将以 CSV 格式接收输入数据。使用 object_detector.DataLoader.from_csv 方法加载数据集并将它们拆分为训练、验证和测试图像。

• 训练图像：这些图像用于训练对象检测模型以识别沙拉配料。
• 验证图像：这些是模型在训练过程中没有见过的图像。您将使用它们来决定何时停止训练，以避免 过拟合。
• 测试图像：这些图像用于评估最终模型的性能。

您可以直接从 Google Cloud Storage 加载 CSV 文件，但您无需将图像保存在 Google Cloud 上即可使用 Model Maker。您可以在计算机上指定本地 CSV 文件，Model Maker 将正常工作。

train_data, validation_data, test_data = object_detector.DataLoader.from_csv('gs://cloud-ml-data/img/openimage/csv/salads_ml_use.csv')

步骤 3. 使用训练数据训练 TensorFlow 模型。

• EfficientDet-Lite0 模型默认使用 epochs = 50，这意味着它将遍历训练数据集 50 次。您可以查看训练期间的验证精度，并尽早停止以避免过拟合。
• 在此处设置 batch_size = 8，这样您将看到遍历训练数据集中的 175 张图像需要 21 个步骤。
• 设置 train_whole_model=True 以微调整个模型，而不仅仅是训练头部层以提高精度。权衡是训练模型可能需要更长时间。

model = object_detector.create(train_data, model_spec=spec, batch_size=8, train_whole_model=True, validation_data=validation_data)

步骤 4. 使用测试数据评估模型。

使用训练数据集中的图像训练对象检测模型后，使用测试数据集中剩余的 25 张图像来评估模型对从未见过的全新数据的性能。

由于默认批次大小为 64，因此遍历测试数据集中 25 张图像需要 1 个步骤。

评估指标与 COCO 相同。

model.evaluate(test_data)

步骤 5. 导出为 TensorFlow Lite 模型。

通过指定要将量化模型导出到的文件夹，将训练后的对象检测模型导出为 TensorFlow Lite 格式。默认的训练后量化技术是全整数量化。

model.export(export_dir='.')

步骤 6. 评估 TensorFlow Lite 模型。

几个因素会影响导出到 TFLite 时模型的精度。

• 量化 有助于将模型大小缩小 4 倍，但会牺牲一些精度。
• 原始 TensorFlow 模型使用每类 非最大抑制 (NMS) 进行后处理，而 TFLite 模型使用全局 NMS，速度快得多，但精度较低。Keras 输出最多 100 个检测结果，而 tflite 输出最多 25 个检测结果。

因此，您需要评估导出的 TFLite 模型，并将其精度与原始 TensorFlow 模型进行比较。

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

您可以使用 Colab 的左侧边栏下载 TensorFlow Lite 模型文件。右键单击 model.tflite 文件，然后选择 下载 以将其下载到本地计算机。

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

有关如何在工作应用程序中使用模型的更多详细信息，请参阅 TFLite 对象检测示例应用程序。

（可选）在您的图像上测试 TFLite 模型

您可以使用来自互联网的图像来测试训练后的 TFLite 模型。

• 将下面的 INPUT_IMAGE_URL 替换为您想要的输入图像。
• 调整 DETECTION_THRESHOLD 以更改模型的灵敏度。较低的阈值意味着模型将拾取更多对象，但也会有更多误检。同时，较高的阈值意味着模型只会拾取它自信地检测到的对象。

尽管目前需要一些样板代码才能在 Python 中运行模型，但将模型集成到移动应用程序中只需要几行代码。

加载训练后的 TFLite 模型并定义一些可视化函数

import cv2

from PIL import Image

model_path = 'model.tflite'

# Load the labels into a list
classes = ['???'] * model.model_spec.config.num_classes
label_map = model.model_spec.config.label_map
for label_id, label_name in label_map.as_dict().items():
  classes[label_id-1] = label_name

# Define a list of colors for visualization
COLORS = np.random.randint(0, 255, size=(len(classes), 3), dtype=np.uint8)

def preprocess_image(image_path, input_size):
  """Preprocess the input image to feed to the TFLite model"""
  img = tf.io.read_file(image_path)
  img = tf.io.decode_image(img, channels=3)
  img = tf.image.convert_image_dtype(img, tf.uint8)
  original_image = img
  resized_img = tf.image.resize(img, input_size)
  resized_img = resized_img[tf.newaxis, :]
  resized_img = tf.cast(resized_img, dtype=tf.uint8)
  return resized_img, original_image


def detect_objects(interpreter, image, threshold):
  """Returns a list of detection results, each a dictionary of object info."""

  signature_fn = interpreter.get_signature_runner()

  # Feed the input image to the model
  output = signature_fn(images=image)

  # Get all outputs from the model
  count = int(np.squeeze(output['output_0']))
  scores = np.squeeze(output['output_1'])
  classes = np.squeeze(output['output_2'])
  boxes = np.squeeze(output['output_3'])

  results = []
  for i in range(count):
    if scores[i] >= threshold:
      result = {
        'bounding_box': boxes[i],
        'class_id': classes[i],
        'score': scores[i]
      }
      results.append(result)
  return results


def run_odt_and_draw_results(image_path, interpreter, threshold=0.5):
  """Run object detection on the input image and draw the detection results"""
  # Load the input shape required by the model
  _, input_height, input_width, _ = interpreter.get_input_details()[0]['shape']

  # Load the input image and preprocess it
  preprocessed_image, original_image = preprocess_image(
      image_path,
      (input_height, input_width)
    )

  # Run object detection on the input image
  results = detect_objects(interpreter, preprocessed_image, threshold=threshold)

  # Plot the detection results on the input image
  original_image_np = original_image.numpy().astype(np.uint8)
  for obj in results:
    # Convert the object bounding box from relative coordinates to absolute
    # coordinates based on the original image resolution
    ymin, xmin, ymax, xmax = obj['bounding_box']
    xmin = int(xmin * original_image_np.shape[1])
    xmax = int(xmax * original_image_np.shape[1])
    ymin = int(ymin * original_image_np.shape[0])
    ymax = int(ymax * original_image_np.shape[0])

    # Find the class index of the current object
    class_id = int(obj['class_id'])

    # Draw the bounding box and label on the image
    color = [int(c) for c in COLORS[class_id]]
    cv2.rectangle(original_image_np, (xmin, ymin), (xmax, ymax), color, 2)
    # Make adjustments to make the label visible for all objects
    y = ymin - 15 if ymin - 15 > 15 else ymin + 15
    label = "{}: {:.0f}%".format(classes[class_id], obj['score'] * 100)
    cv2.putText(original_image_np, label, (xmin, y),
        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

  # Return the final image
  original_uint8 = original_image_np.astype(np.uint8)
  return original_uint8

运行对象检测并显示检测结果

INPUT_IMAGE_URL = "https://storage.googleapis.com/cloud-ml-data/img/openimage/3/2520/3916261642_0a504acd60_o.jpg"
DETECTION_THRESHOLD = 0.3

TEMP_FILE = '/tmp/image.png'

!wget -q -O $TEMP_FILE $INPUT_IMAGE_URL
im = Image.open(TEMP_FILE)
im.thumbnail((512, 512), Image.ANTIALIAS)
im.save(TEMP_FILE, 'PNG')

# Load the TFLite model
interpreter = tf.lite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()

# Run inference and draw detection result on the local copy of the original file
detection_result_image = run_odt_and_draw_results(
    TEMP_FILE,
    interpreter,
    threshold=DETECTION_THRESHOLD
)

# Show the detection result
Image.fromarray(detection_result_image)

（可选）为 Edge TPU 编译

现在您已经拥有了一个量化的 EfficientDet Lite 模型，可以将其编译并部署到 Coral EdgeTPU。

步骤 1. 安装 EdgeTPU 编译器

 curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -

 echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list

 sudo apt-get update

 sudo apt-get install edgetpu-compiler

步骤 2. 选择 Edge TPU 的数量，编译

EdgeTPU 拥有 8MB 的 SRAM 用于缓存模型参数（更多信息）。这意味着对于大于 8MB 的模型，推理时间将增加，以便传输模型参数。避免这种情况的一种方法是 模型流水线 - 将模型拆分为可以拥有专用 EdgeTPU 的片段。这可以显着提高延迟。

下表可作为使用 Edge TPU 数量的参考 - 较大的模型将无法为单个 TPU 编译，因为中间张量无法容纳在片上内存中。

NUMBER_OF_TPUS =  1

!edgetpu_compiler model.tflite --num_segments=$NUMBER_OF_TPUS

步骤 3. 下载，运行模型

编译完模型后，现在可以在 EdgeTPU 上运行它们以进行对象检测。首先，使用 Colab 的左侧边栏下载编译后的 TensorFlow Lite 模型文件。右键单击 model_edgetpu.tflite 文件，然后选择 下载 以将其下载到本地计算机。

现在，您可以以您喜欢的方式运行模型。检测示例包括

• pycoral 检测
• 基本 TFLite 检测
• 示例视频检测
• libcoral C++ API

高级用法

本节介绍调整模型和训练超参数等高级用法主题。

加载数据集

加载您自己的数据

您可以上传您自己的数据集以完成本教程。使用 Colab 中的左侧边栏上传您的数据集。

如果您不想将数据集上传到云端，也可以按照 指南 在本地运行库。

使用不同的数据格式加载您的数据

Model Maker 库还支持 object_detector.DataLoader.from_pascal_voc 方法来加载具有 PASCAL VOC 格式的数据。 makesense.ai 和 LabelImg 是可以注释图像并将注释保存为 PASCAL VOC 数据格式的 XML 文件的工具。

object_detector.DataLoader.from_pascal_voc(image_dir, annotations_dir, label_map={1: "person", 2: "notperson"})

自定义 EfficientDet 模型超参数

您可以调整的模型和训练管道参数如下

• model_dir：保存模型检查点文件的路径。如果未设置，将使用临时目录。
• steps_per_execution：每次训练执行的步骤数。
• moving_average_decay：浮点数。用于维护训练参数的移动平均值的衰减。
• var_freeze_expr：正则表达式，用于映射要冻结的变量的前缀名称，这意味着在训练期间保持不变。更具体地说，在代码库中使用 re.match(var_freeze_expr, variable_name) 来映射要冻结的变量。
• tflite_max_detections：整数，默认值为 25。TFLite 模型中输出检测结果的最大数量。
• strategy：指定要使用的分布式策略的字符串。接受的值为“tpu”、“gpus”、None。“tpu” 表示使用 TPUStrategy。“gpus” 表示对多个 GPU 使用 MirroredStrategy。如果为 None，则使用 TF 默认值，即 OneDeviceStrategy。
• tpu：用于训练的 Cloud TPU。这应该是创建 Cloud TPU 时使用的名称，或者是一个 grpc://ip.address.of.tpu:8470 URL。
• use_xla：即使策略不是 tpu，也要使用 XLA。如果策略是 tpu，则始终使用 XLA，此标志无效。
• profile：启用配置文件模式。
• debug：启用调试模式。

其他可调整的参数在 hparams_config.py 中显示。

例如，您可以设置 var_freeze_expr='efficientnet'，它会冻结名称前缀为 efficientnet 的变量（默认值为 '(efficientnet|fpn_cells|resample_p6)'）。这允许模型冻结不可训练的变量，并在训练过程中保持其值不变。

spec = model_spec.get('efficientdet_lite0')
spec.config.var_freeze_expr = 'efficientnet'

更改模型架构

您可以通过更改 model_spec 来更改模型架构。例如，将 model_spec 更改为 EfficientDet-Lite4 模型。

spec = model_spec.get('efficientdet_lite4')

调整训练超参数

The create 函数是 Model Maker 库用来创建模型的驱动函数。The model_spec 参数定义了模型规范。The object_detector.EfficientDetSpec 类目前受支持。The create 函数包含以下步骤

1. 根据 model_spec 创建用于目标检测的模型。
2. 训练模型。默认的 epoch 和默认的批次大小由 model_spec 对象中的 epochs 和 batch_size 变量设置。您还可以调整影响模型精度的训练超参数，例如 epochs 和 batch_size。例如，

• epochs：整数，默认值为 50。更多的 epoch 可能可以获得更高的精度，但可能会导致过拟合。
• batch_size：整数，默认值为 64。在一个训练步骤中使用的样本数量。
• train_whole_model：布尔值，默认值为 False。如果为 True，则训练整个模型。否则，只训练与 var_freeze_expr 不匹配的层。

例如，您可以使用较少的 epoch 训练，并且只训练头部层。您可以增加 epoch 数以获得更好的结果。

model = object_detector.create(train_data, model_spec=spec, epochs=10, validation_data=validation_data)

导出到不同的格式

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

• ExportFormat.TFLITE
• ExportFormat.LABEL
• ExportFormat.SAVED_MODEL

默认情况下，它只导出包含模型 元数据 的 TensorFlow Lite 模型文件，以便您可以在以后的设备上使用 ML 应用程序。标签文件嵌入在元数据中。

在许多设备上的 ML 应用程序中，模型大小是一个重要的因素。因此，建议您量化模型以使其更小，并可能运行得更快。对于 EfficientDet-Lite 模型，默认情况下使用全整数量化来量化模型。有关更多详细信息，请参阅 训练后量化。

model.export(export_dir='.')

您也可以选择导出与模型相关的其他文件以进行更好的检查。例如，导出保存的模型和标签文件，如下所示

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

自定义 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)

阅读更多

您可以阅读我们的 目标检测 示例以了解技术细节。有关更多信息，请参阅

• TensorFlow Lite Model Maker 指南 和 API 参考。
• 任务库：ObjectDetector 用于部署。
• 端到端参考应用程序：Android、iOS 和 Raspberry PI。