使用 TensorFlow Lite Model Maker 的 BERT 问答

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

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_negativeFalse

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_epochsdefault_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 模型文件。

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

默认情况下,它仅导出带有元数据的 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_lenquery_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

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

调整训练超参数

您还可以调整训练超参数(例如 epochsbatch_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 问答示例以了解技术细节。如需了解更多信息,请参阅: