使用 TensorFlow Lite Model Maker 进行文本分类

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

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

此笔记本展示了一个端到端示例，该示例利用 Model Maker 库来说明如何调整和转换常用的文本分类模型，以便在移动设备上对电影评论进行分类。文本分类模型将文本分类为预定义的类别。输入应该是预处理的文本，输出是类别的概率。本教程中使用的数据集是正面和负面的电影评论。

先决条件

安装所需的软件包

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

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

导入所需的软件包。

import numpy as np
import os

from tflite_model_maker import model_spec
from tflite_model_maker import text_classifier
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.text_classifier import AverageWordVecSpec
from tflite_model_maker.text_classifier import DataLoader

from tflite_support.task import core
from tflite_support.task import processor
from tflite_support.task import text

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

下载示例训练数据。

在本教程中，我们将使用 SST-2（斯坦福情感树库），它是 GLUE 基准测试中的任务之一。它包含 67,349 条电影评论用于训练，以及 872 条电影评论用于测试。数据集有两个类别：正面和负面电影评论。

data_dir = tf.keras.utils.get_file(
      fname='SST-2.zip',
      origin='https://dl.fbaipublicfiles.com/glue/data/SST-2.zip',
      extract=True)
data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')

SST-2 数据集以 TSV 格式存储。TSV 和 CSV 之间的唯一区别是 TSV 使用制表符 \t 字符作为分隔符，而不是 CSV 格式中的逗号 ,。

以下是训练数据集的前 5 行。label=0 表示负面，label=1 表示正面。

句子 标签
隐藏来自父母单位的新分泌物 0
不含任何智慧，只有费力的笑料 0
它热爱自己的角色，并传达了关于人性的某种美丽 1
始终保持满足，始终保持不变 0
关于电影制作人所能挖出的最糟糕的“书呆子复仇”陈词滥调 0

接下来，我们将数据集加载到 Pandas 数据框中，并将当前标签名称（0 和 1）更改为更易于理解的名称（negative 和 positive），并将其用于模型训练。

import pandas as pd

def replace_label(original_file, new_file):
  # Load the original file to pandas. We need to specify the separator as
  # '\t' as the training data is stored in TSV format
  df = pd.read_csv(original_file, sep='\t')

  # Define how we want to change the label name
  label_map = {0: 'negative', 1: 'positive'}

  # Excute the label change
  df.replace({'label': label_map}, inplace=True)

  # Write the updated dataset to a new file
  df.to_csv(new_file)

# Replace the label name for both the training and test dataset. Then write the
# updated CSV dataset to the current folder.
replace_label(os.path.join(os.path.join(data_dir, 'train.tsv')), 'train.csv')
replace_label(os.path.join(os.path.join(data_dir, 'dev.tsv')), 'dev.csv')

快速入门

训练文本分类模型有五个步骤

步骤 1. 选择文本分类模型架构。

这里我们使用平均词嵌入模型架构，它将生成一个体积小且速度快的模型，并具有不错的准确率。

spec = model_spec.get('average_word_vec')

Model Maker 还支持其他模型架构，例如 BERT。如果您有兴趣了解其他架构，请参阅下面的 为文本分类器选择模型架构 部分。

步骤 2. 加载训练和测试数据，然后根据特定的 model_spec 对其进行预处理。

Model Maker 可以接受 CSV 格式的输入数据。我们将使用之前创建的具有易于理解的标签名称的训练和测试数据集进行加载。

每个模型架构都需要以特定方式处理输入数据。 DataLoader 从 model_spec 中读取要求，并自动执行必要的预处理。

train_data = DataLoader.from_csv(
      filename='train.csv',
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      is_training=True)
test_data = DataLoader.from_csv(
      filename='dev.csv',
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      is_training=False)

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

平均词嵌入模型默认使用 batch_size = 32。因此，您会看到它需要 2104 步才能遍历训练数据集中 67,349 个句子。我们将训练模型 10 个 epochs，这意味着遍历训练数据集 10 次。

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

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

在使用训练数据集中句子训练文本分类模型后，我们将使用测试数据集中剩余的 872 个句子来评估模型在从未见过的新数据上的表现。

由于默认批次大小为 32，因此遍历测试数据集中 872 个句子需要 28 步。

loss, acc = model.evaluate(test_data)

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

让我们将我们训练的文本分类模型导出为 TensorFlow Lite 格式。我们将指定要导出模型的文件夹。默认情况下，浮点 TFLite 模型将为平均词嵌入模型架构导出。

model.export(export_dir='average_word_vec')

您可以使用 Colab 的左侧边栏下载 TensorFlow Lite 模型文件。进入 average_word_vec 文件夹（如我们在上面的 export_dir 参数中指定的），右键单击 model.tflite 文件，然后选择 Download 将其下载到本地计算机。

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

有关如何在工作应用程序中使用模型的更多详细信息，请参阅 TFLite 文本分类示例应用程序。

注意 1：Android Studio 模型绑定尚不支持文本分类，因此请使用 TensorFlow Lite 任务库。

注意 2：TFLite 模型所在的同一文件夹中有一个 model.json 文件。它包含捆绑在 TensorFlow Lite 模型中的 元数据 的 JSON 表示形式。模型元数据有助于 TFLite 任务库了解模型的功能以及如何对模型进行数据预处理/后处理。您无需下载 model.json 文件，因为它仅用于信息目的，其内容已包含在 TFLite 文件中。

注意 3：如果您使用 MobileBERT 或 BERT-Base 架构训练文本分类模型，则需要使用 BertNLClassifier API 来将训练后的模型集成到移动应用程序中。

以下部分将逐步介绍示例，以显示更多详细信息。

步骤 6：使用 TFLite 任务库 演示如何使用训练后的模型

将 dev.csv 文件读入句子数据，以便使用训练后的模型进行预测

sentence_data = pd.read_csv('/content/dev.csv', index_col=0)
sentence_data

模型配置参数

# Name of the TFLite text classification model.
_MODEL = '/content/average_word_vec/model.tflite'
# Whether to run the model on EdgeTPU.
_ENABLE_EDGETPU = False
# Number of CPU threads to run the model.
_NUM_THREADS = 4

初始化模型

我们还可以更改可能影响模型结果的参数，例如 file_name、use_coral 和 num_threads。您可以调整的参数是

  • file_name：TFLite 图像分类模型的名称。
  • use_coral：如果为 true，则推理将委托给连接的 Coral Edge TPU 设备。
  • num_threads：用于运行模型的 CPU 线程数。
# Initialize the text classification model.
base_options = core.BaseOptions(file_name=_MODEL, use_coral=_ENABLE_EDGETPU, num_threads=_NUM_THREADS)
options = text.NLClassifierOptions(base_options)

# Create NLClassifier from options.
classifier = text.NLClassifier.create_from_options(options)

使用 TFLite 任务库 进行预测

for idx in range(20):
  sentence = sentence_data['sentence'].iloc[idx]
  label = sentence_data['label'].iloc[idx]
  text_classification_result = classifier.classify(sentence)
  classification_list = text_classification_result.classifications[0].categories

  # Sort output by probability descending.
  predict_label = sorted(
      classification_list, key=lambda item: item.score, reverse=True)[0]

  print('truth_label: {} -----> predict_label: {}'.format(label, predict_label.category_name))

为文本分类器选择模型架构

每个 model_spec 对象代表文本分类器的特定模型。TensorFlow Lite 模型制作器目前支持 MobileBERT、平均词嵌入和 BERT-Base 模型。

支持的模型 model_spec 的名称 模型描述 模型大小
平均词嵌入 'average_word_vec' 使用 RELU 激活函数对文本词嵌入进行平均。 <1MB
MobileBERT 'mobilebert_classifier' 比 BERT-Base 小 4.3 倍，快 5.5 倍，同时取得了具有竞争力的结果，适合于设备上应用。 25MB（带量化）
100MB（不带量化）
BERT-Base 'bert_classifier' 广泛用于 NLP 任务的标准 BERT 模型。 300MB

在快速入门中，我们使用了平均词嵌入模型。让我们切换到 MobileBERT，以训练具有更高精度的模型。

mb_spec = model_spec.get('mobilebert_classifier')

加载训练数据

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

Upload File

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

为了简单起见，我们将重复使用之前下载的 SST-2 数据集。让我们使用 DataLoader.from_csv 方法加载数据。

请注意，由于我们更改了模型架构，因此需要重新加载训练和测试数据集以应用新的预处理逻辑。

train_data = DataLoader.from_csv(
      filename='train.csv',
      text_column='sentence',
      label_column='label',
      model_spec=mb_spec,
      is_training=True)
test_data = DataLoader.from_csv(
      filename='dev.csv',
      text_column='sentence',
      label_column='label',
      model_spec=mb_spec,
      is_training=False)

模型制作器库还支持 from_folder() 方法来加载数据。它假设同一类的文本数据位于同一子目录中，并且子文件夹名称是类名。每个文本文件包含一个电影评论样本。 class_labels 参数用于指定子文件夹。

训练 TensorFlow 模型

使用训练数据训练文本分类模型。

model = text_classifier.create(train_data, model_spec=mb_spec, epochs=3)

检查详细的模型结构。

model.summary()

评估模型

使用测试数据评估我们刚刚训练的模型，并测量损失和准确率值。

loss, acc = model.evaluate(test_data)

导出为 TensorFlow Lite 模型

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

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

model.export(export_dir='mobilebert/')

TensorFlow Lite 模型文件可以使用 BertNLClassifier API（位于 TensorFlow Lite 任务库 中）集成到移动应用程序中。请注意，这与用于集成使用平均词向量模型架构训练的文本分类的 NLClassifier API 不同。

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

默认情况下，它只导出包含模型元数据的 TensorFlow Lite 模型文件。您也可以选择导出与模型相关的其他文件，以便更好地进行检查。例如，仅导出标签文件和词汇表文件，如下所示

model.export(export_dir='mobilebert/', export_format=[ExportFormat.LABEL, ExportFormat.VOCAB])

您可以使用 evaluate_tflite 方法评估 TFLite 模型，以测量其准确率。将训练后的 TensorFlow 模型转换为 TFLite 格式并应用量化可能会影响其准确率，因此建议在部署之前评估 TFLite 模型的准确率。

accuracy = model.evaluate_tflite('mobilebert/model.tflite', test_data)
print('TFLite model accuracy: ', accuracy)

高级用法

create 函数是模型制作器库用于创建模型的驱动函数。 model_spec 参数定义模型规范。目前支持 AverageWordVecSpec 和 BertClassifierSpec 类。 create 函数包含以下步骤

  1. 根据 model_spec 创建文本分类器的模型。
  2. 训练分类器模型。默认的 epochs 和默认的批次大小由 model_spec 对象中的 default_training_epochs 和 default_batch_size 变量设置。

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

自定义 MobileBERT 模型超参数

您可以调整的模型参数是

  • seq_len：馈送到模型的序列的长度。
  • initializer_range：用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • trainable：布尔值，指定预训练层是否可训练。

您可以调整的训练管道参数是

  • model_dir：模型检查点文件的存储位置。如果未设置，将使用临时目录。
  • dropout_rate：dropout 率。
  • learning_rate：Adam 优化器的初始学习率。
  • tpu：要连接的 TPU 地址。

例如，您可以将 seq_len=256（默认值为 128）。这允许模型对更长的文本进行分类。

new_model_spec = model_spec.get('mobilebert_classifier')
new_model_spec.seq_len = 256

自定义平均词嵌入模型超参数

您可以调整模型基础结构，例如 AverageWordVecSpec 类中的 wordvec_dim 和 seq_len 变量。

例如，您可以使用更大的 wordvec_dim 值训练模型。请注意，如果您修改了模型，则必须构造一个新的 model_spec。

new_model_spec = AverageWordVecSpec(wordvec_dim=32)

获取预处理后的数据。

new_train_data = DataLoader.from_csv(
      filename='train.csv',
      text_column='sentence',
      label_column='label',
      model_spec=new_model_spec,
      is_training=True)

训练新模型。

model = text_classifier.create(new_train_data, model_spec=new_model_spec)

调整训练超参数

您还可以调整训练超参数，例如 epochs 和 batch_size，它们会影响模型的准确率。例如，

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

例如，您可以使用更多 epochs 进行训练。

model = text_classifier.create(new_train_data, model_spec=new_model_spec, epochs=20)

使用 20 个训练 epochs 评估新训练的模型。

new_test_data = DataLoader.from_csv(
      filename='dev.csv',
      text_column='sentence',
      label_column='label',
      model_spec=new_model_spec,
      is_training=False)

loss, accuracy = model.evaluate(new_test_data)

更改模型架构

您可以通过更改 model_spec 来更改模型。以下展示了如何更改为 BERT-Base 模型。

将 model_spec 更改为 BERT-Base 模型，用于文本分类器。

spec = model_spec.get('bert_classifier')

其余步骤相同。

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

了解更多

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