在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

TensorFlow Lite Model Maker 库可以简化将 TensorFlow 模型调整和转换为特定输入数据的过程，以便在部署用于设备上机器学习应用程序时使用该模型。

此笔记本展示了一个端到端示例，该示例使用 Model Maker 库来说明如何调整和转换常用的问答模型以用于问答任务。

BERT 问答任务简介

此库支持的 task 是抽取式问答 task，这意味着给定一段文字和一个问题，答案是该段文字中的一个片段。下图显示了一个问答示例。

答案是段落中的片段（图片来自：SQuAD 博客）

对于问答 task 的模型，输入应该是已经预处理过的段落和问题对，输出应该是段落中每个 token 的起始 logits 和结束 logits。输入的大小可以根据段落和问题的长度进行设置和调整。

端到端概述

以下代码片段演示了如何在几行代码内获取模型。整个过程包括 5 个步骤：（1）选择模型，（2）加载数据，（3）重新训练模型，（4）评估，以及（5）将其导出为 TensorFlow Lite 格式。

# Chooses a model specification that represents the model.
spec = model_spec.get('mobilebert_qa')

# Gets the training data and validation data.
train_data = DataLoader.from_squad(train_data_path, spec, is_training=True)
validation_data = DataLoader.from_squad(validation_data_path, spec, is_training=False)

# Fine-tunes the model.
model = question_answer.create(train_data, model_spec=spec)

# Gets the evaluation result.
metric = model.evaluate(validation_data)

# Exports the model to the TensorFlow Lite format with metadata in the export directory.
model.export(export_dir)

以下部分将更详细地解释代码。

先决条件

要运行此示例，请安装所需的软件包，包括来自 GitHub 仓库 的 Model Maker 软件包。

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

导入所需的软件包。

import numpy as np
import os

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

from tflite_model_maker import model_spec
from tflite_model_maker import question_answer
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.question_answer import DataLoader

“端到端概述”演示了一个简单的端到端示例。以下部分将逐步介绍该示例，以显示更多详细信息。

选择一个表示问答模型的 model_spec

每个 model_spec 对象都表示一个特定的问答模型。Model Maker 目前支持 MobileBERT 和 BERT-Base 模型。

支持的模型	model_spec 的名称	模型描述
MobileBERT	'mobilebert_qa'	比 BERT-Base 小 4.3 倍，快 5.5 倍，同时实现了具有竞争力的结果，适用于设备上场景。
MobileBERT-SQuAD	'mobilebert_qa_squad'	与 MobileBERT 模型具有相同的模型架构，并且初始模型已在 SQuAD1.1 上重新训练。
BERT-Base	'bert_qa'	广泛用于 NLP 任务的标准 BERT 模型。

在本教程中，MobileBERT-SQuAD 用作示例。由于该模型已在 SQuAD1.1 上重新训练，因此它可以更快地完成问答 task。

spec = model_spec.get('mobilebert_qa_squad')

加载特定于设备上机器学习应用程序的输入数据并预处理数据

TriviaQA 是一个阅读理解数据集，包含超过 650K 个问答证据三元组。在本教程中，您将使用此数据集的一个子集来学习如何使用 Model Maker 库。

要加载数据，请通过运行 转换 Python 脚本 将 TriviaQA 数据集转换为 SQuAD1.1 格式，并使用 --sample_size=8000 和一组 web 数据。通过以下方式稍微修改转换代码：

• 跳过在上下文文档中找不到任何答案的样本；
• 获取上下文中原始答案，不区分大小写。

下载已转换数据集的存档版本。

train_data_path = tf.keras.utils.get_file(
    fname='triviaqa-web-train-8000.json',
    origin='https://storage.googleapis.com/download.tensorflow.org/models/tflite/dataset/triviaqa-web-train-8000.json')
validation_data_path = tf.keras.utils.get_file(
    fname='triviaqa-verified-web-dev.json',
    origin='https://storage.googleapis.com/download.tensorflow.org/models/tflite/dataset/triviaqa-verified-web-dev.json')

您也可以使用自己的数据集训练 MobileBERT 模型。如果您在 Colab 上运行此笔记本，请使用左侧边栏上传您的数据。

如果您不想将数据上传到云端，也可以按照 指南 在离线状态下运行库。

使用 DataLoader.from_squad 方法根据特定的 model_spec 加载和预处理 SQuAD 格式 数据。您可以使用 SQuAD2.0 或 SQuAD1.1 格式。将参数 version_2_with_negative 设置为 True 表示格式为 SQuAD2.0。否则，格式为 SQuAD1.1。默认情况下，version_2_with_negative 为 False。

train_data = DataLoader.from_squad(train_data_path, spec, is_training=True)
validation_data = DataLoader.from_squad(validation_data_path, spec, is_training=False)

自定义 TensorFlow 模型

基于加载的数据创建自定义问答模型。 create 函数包含以下步骤

1. 根据 model_spec 创建问答模型。
2. 训练问答模型。默认的 epoch 和默认的批次大小根据 model_spec 对象中的两个变量 default_training_epochs 和 default_batch_size 设置。

model = question_answer.create(train_data, model_spec=spec)

查看详细的模型结构。

model.summary()

评估自定义模型

在验证数据上评估模型，并获取包含 f1 分数和 exact match 等指标的字典。请注意，SQuAD1.1 和 SQuAD2.0 的指标不同。

model.evaluate(validation_data)

导出到 TensorFlow Lite 模型

将训练后的模型转换为 TensorFlow Lite 模型格式，并包含 元数据，以便您以后在设备上运行的 ML 应用程序中使用。词汇文件嵌入在元数据中。默认的 TFLite 文件名为 model.tflite。

在许多设备上运行的 ML 应用程序中，模型大小是一个重要因素。因此，建议您对模型进行量化，以使其更小，并可能更快地运行。BERT 和 MobileBERT 模型的默认后训练量化技术是动态范围量化。

model.export(export_dir='.')

您可以在 bert_qa 参考应用程序中使用 TensorFlow Lite 模型文件，使用 BertQuestionAnswerer API 在 TensorFlow Lite 任务库 中使用，您可以从 Colab 上的左侧边栏下载它。

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

• ExportFormat.TFLITE
• ExportFormat.VOCAB
• ExportFormat.SAVED_MODEL

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

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

您也可以使用 evaluate_tflite 方法评估 tflite 模型。此步骤预计需要很长时间。

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

高级用法

create 函数是此库的关键部分，其中 model_spec 参数定义模型规范。目前支持 BertQASpec 类。有 2 个模型：MobileBERT 模型、BERT-Base 模型。 create 函数包含以下步骤

1. 根据 model_spec 创建问答模型。
2. 训练问答模型。

本节介绍了几个高级主题，包括调整模型、调整训练超参数等。

调整模型

您可以调整模型基础架构，例如 BertQASpec 类中的参数 seq_len 和 query_len。

模型的可调整参数

• seq_len：输入模型的段落的长度。
• query_len：输入模型的问题的长度。
• doc_stride：使用滑动窗口方法获取文档块时的步长。
• initializer_range：用于初始化所有权重矩阵的截断正态分布初始化器的标准差。
• trainable：布尔值，表示预训练层是否可训练。

训练管道的可调整参数

• model_dir：模型检查点文件的位置。如果未设置，将使用临时目录。
• dropout_rate：dropout 的比率。
• learning_rate：Adam 的初始学习率。
• predict_batch_size：预测的批次大小。
• tpu：要连接的 TPU 地址。仅在使用 tpu 时使用。

例如，您可以使用更长的序列长度训练模型。如果您更改模型，则必须首先构造一个新的 model_spec。

new_spec = model_spec.get('mobilebert_qa')
new_spec.seq_len = 512

其余步骤相同。请注意，您必须重新运行 dataloader 和 create 部分，因为不同的模型规范可能具有不同的预处理步骤。

调整训练超参数

您还可以调整训练超参数，例如 epochs 和 batch_size，以影响模型性能。例如，

• epochs：更多的 epoch 可以获得更好的性能，但可能会导致过拟合。
• batch_size：一次训练步骤中使用的样本数量。

例如，您可以使用更多 epoch 和更大的批次大小进行训练，例如

model = question_answer.create(train_data, model_spec=spec, epochs=5, batch_size=64)

更改模型架构

您可以通过更改 model_spec 来更改数据训练的基础模型。例如，要更改为 BERT-Base 模型，请运行

spec = model_spec.get('bert_qa')

其余步骤相同。

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

阅读更多

您可以阅读我们的 BERT 问答 示例以了解技术细节。有关更多信息，请参阅

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