使用 TensorFlow Lite Model Maker 的 BERT 问答

TensorFlow Lite Model Maker 库简化了在设备端机器学习应用中部署 TensorFlow 模型时修改此模型并将其转换为特定输入数据的过程。

此笔记本展示了一个端到端示例，该示例利用 Model Maker 库说明了如何对用于问答任务的常用问答模型进行修改和转换。

BERT 问答任务简介

此库中支持的任务是提取性问答任务，这意味着给定一个段落和一个问题，答案就是段落中的跨度。下图给出了一个问答示例。

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

对于问答任务的模型，输入应当是经过预处理的段落和问题对，输出应当是段落中每个词例的起始 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 进行了重新训练，因此可以更快地覆盖问答任务。

spec = model_spec.get('mobilebert_qa_squad')

加载特定于设备端 ML 应用的输入数据并对这些数据进行预处理

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

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

• 跳过无法在上下文文档中找到任何答案的样本；
• 在不使用大写或小写的上下文中获取原始答案。

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

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. 训练问答模型。根据 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='.')

您可以通过从 Colab 的左侧边栏中下载并使用 TensorFlow Lite Task Library 中的 BertQuestionAnswerer API，在 bert_qa 参考应用中使用 TensorFlow Lite 模型文件。

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

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

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

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：truncated_normal_initializer 的 stdev，用于初始化所有权重矩阵。
• trainable：布尔值，预训练层是否可训练。

训练流水线的可调参数：

• model_dir：模型检查点文件的位置。如果未设置，将使用临时目录。
• dropout_rate：随机失活率。
• 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：更多周期可能会获得更好的性能，但也可能导致过拟合。
• batch_size：一个训练步骤中要使用的样本数。

例如，您可以使用更多的周期和更大的批次大小进行训练，代码如下：

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 参考。
• Task Library：用于部署的 ImageClassifier。
• 端到端参考应用：Android 和 iOS。