作者：Tian Lin、Yicheng Fan、Jaesung Chung 和 Chen Cen

TensorFlow Lite 已广泛应用于许多应用程序中，为移动电话、微控制器单元和边缘 TPU 等边缘设备提供机器学习功能。在所有让生活更轻松、更高效的热门应用中，自然语言理解是研究界和行业都高度关注的关键领域之一。在 2019 年 TensorFlow World 上展示了设备上问答用例的 演示 之后，我们收到了社区许多关于如何让更多此类 NLP 模型可用于设备上推理的反馈。

在这些反馈的启发下，我们很高兴地宣布，TensorFlow Lite 现在全面支持 NLP 任务。有了这项基础设施工作，越来越多的 NLP 模型能够在移动电话上运行，用户可以享受 NLP 模型的优势，同时将个人数据保留在设备上。在本博文中，我们将介绍允许执行以下操作的新功能：（1）使用新的预训练 NLP 模型，（2）创建自己的 NLP 模型，（3）更好地支持将 TensorFlow NLP 模型转换为 TensorFlow Lite 格式，以及（4）将这些模型部署到移动设备上。

使用新的预训练 NLP 模型

参考应用

参考应用是一组开源 移动应用程序，它们封装了预训练的机器学习 模型、推理代码和可运行的演示。我们提供了一系列与 Android Studio 和 XCode 集成的 NLP 参考应用，因此开发人员只需单击一下即可构建，然后部署到 Android 或 iOS 手机上。

使用以下 NLP 参考应用，移动开发人员可以学习将现有 NLP 模型（由 BERT、MobileBERT 或 ALBERT 提供支持）集成在一起的端到端流程，转换原始文本数据，并将模型的输入和输出连接起来以生成预测结果，

• 文本分类：模型根据给定的文本数据预测标签。
  
• 问答应用程序：给定一篇文章和用户的问题，模型可以在文章中回答问题。
  
• 智能回复应用程序：给定之前的上下文，模型可以预测并生成潜在的自动回复。
  

上述参考应用中使用的预训练模型可在 TensorFlow Hub 中获取。下图显示了模型之间延迟、大小和 F1 分数的比较。

Pixel 4 CPU、4 个线程上的基准测试，2020 年 3 月 模型超参数：序列长度 128，词典大小 30K

针对设备上用例优化 NLP 模型

与服务器端模型相比，设备上模型具有不同的限制。它们在内存较少且芯片速度较慢的设备上运行，因此需要针对大小和推理速度进行优化。以下是一些针对 NLP 任务优化模型的方法示例。

量化 MobileBERT

MobileBERT 是一个紧凑的 BERT 模型，已在 GitHub 上开源。它比 BERT 基础版 (float32) 小 4.3 倍，速度快 5.5 倍，同时在 GLUE 和 SQuAD 数据集上取得了可比的结果。\\

在初始版本发布后，我们通过使用 量化 进一步改进了模型，以优化其模型大小和性能，以便它可以利用 GPU/DSP 等加速器（如果可用）。量化的 MobileBERT 比 BERT 基础版小 16 倍，速度快 8 倍，但精度损失很小。用于移动设备的 MLPerf 社区正在利用量化的 MobileBERT 模型进行移动推理基准测试，该模型还可以使用 TensorFlow.js 在 Chrome 中运行。

与原始 BERT 基础模型 (416MB) 相比，下表显示了在相同设置下量化 MobileBERT 的性能。

使用投影方法的无嵌入 NLP 模型

语言识别是一种对给定文本的语言进行分类的问题类型。最近，我们 开源 了两个使用投影方法的模型，即 SGNN 和 PRADO。

我们使用 SGNN 来展示如何轻松高效地将 Tensorflow Lite 用于 NLP 任务。SGNN 将文本投影到固定长度的特征，然后进行全连接层运算。通过注释来告诉 TensorFlow Lite 转换器融合 TF.Text API，我们可以获得一个更有效的模型，以便在 TensorFlow Lite 上进行推理。以前，该模型在基准测试中需要 1332.87 μs 才能运行；融合之后，我们在同一台机器上看到了 64.06 μs。这带来了神奇的 20 倍加速效果。

我们还演示了一种名为 PRADO 的模型架构。PRADO 首先从单词标记序列中计算可训练的投影特征，然后应用卷积和注意力将特征映射到固定长度的编码。通过结合投影层、卷积和注意力编码器机制，PRADO 实现了与 LSTM 相似的精度，但模型大小却小了 100 倍。

这些模型背后的理念是使用投影从文本中计算特征，这样模型就不需要维护一个大型嵌入表来将文本特征转换为嵌入。通过这种方式，我们已经证明，该模型将比基于嵌入的模型小得多，同时还能保持相似的性能和推理延迟。

创建您自己的 NLP 模型

除了使用预训练模型之外，TensorFlow Lite 还为您提供了 Model Maker 等工具，让您可以根据自己的数据定制现有模型。

TensorFlow Lite Model Maker：面向机器学习初学者的迁移学习工具包

TensorFlow Lite Model Maker 是一款易于使用的迁移学习工具，用于将最先进的机器学习模型调整到您的数据集。它允许移动开发人员在没有任何机器学习专业知识的情况下创建模型，通过迁移学习减少所需的训练数据并缩短训练时间。

在最初的版本专注于视觉任务后，我们最近在 Model Maker 中添加了两个新的 NLP 任务。您可以按照 文本分类 和 问答 的 Colab 和指南进行操作。要安装 Model Maker

pip install tflite-model-maker

要定制模型，开发人员需要编写以下几行 Python 代码

# Loads Data.
train_data = TextClassifierDataLoader.from_csv(train_csv_file, mode_spec=spec)
test_data = TextClassifierDataLoader.from_csv(test_csv_file, mode_spec=spec)

# Customize the TensorFlow model.
model = text_classifier.create(train_data, model_spec=spec)

# Evaluate the model.
loss, acc = model.evaluate(test_data)

# Export as a TensorFlow Lite model.
model.export(export_dir, quantization_config=config)

转换：更好地支持将 NLP 模型转换为 TensorFlow Lite

由于 TensorFlow Lite 内置运算符库仅支持 TensorFlow 运算符的 子集，因此您在将 NLP 模型转换为 TensorFlow Lite 时可能会遇到问题，原因可能是缺少运算符或不支持的数据类型（例如 RaggedTensor 支持、哈希表支持和资产文件处理等）。以下是有关如何在这些情况下解决转换问题的一些提示。

在 TensorFlow Lite 中运行 TensorFlow 运算符和 TF.text 运算符

我们已经增强了 选择 TensorFlow 运算符 以支持各种情况。使用选择 TF 运算符，开发人员可以利用 TensorFlow 运算符在 TensorFlow Lite 上运行模型，而无需使用 TensorFlow Lite 内置的等效运算符。例如，在训练 TensorFlow 模型时，通常使用 TF.Text 运算符和 RaggedTensor，现在这些模型可以轻松转换为 TensorFlow Lite 并在需要时使用必要的运算符进行运行。

此外，我们提供了使用运算符选择性 构建 的解决方案，这样我们就可以获得一个用于移动部署的修剪后的二进制文件。它可以选择最终构建目标中使用的一组较小的运算符，从而减少部署中的二进制文件大小。

更高效且友好的自定义运算符

在 TensorFlow Lite 中，我们提供了一些新的 移动友好型运算符，用于 NLP，例如 Ngram、SentencePieceTokenizer、WordPieceTokenizer 和 WhitespaceTokenizer。

以前，有一些限制阻止使用 SentencePiece 的模型转换为 TensorFlow Lite。用于移动设备的新 SentencePieceTokenizer API 解决了这些挑战，同时还优化了实现，使其运行速度更快。

类似地，Ngram 和 WhitespaceTokenizer 现在不仅受支持，而且还能在设备上更高效地执行。

TensorFlow Lite 最近宣布了与 MLIR 的运算符 融合。我们使用相同的机制将 TF.Text API 融合到自定义 TensorFlow Lite 运算符中，从而显着提高了推理效率。例如，WhitespaceTokenizer API 由多个运算符组成，在 TensorFlow Lite 中的原始图中需要 0.9 毫秒才能运行。将这些运算符融合到单个运算符后，它在 0.04 毫秒内完成，速度提高了 23 倍。这种方法已被证明可以极大地提高上述 SGNN 模型的推理延迟。

哈希表支持

哈希表对于许多 NLP 模型来说非常重要，因为我们通常需要通过将单词转换为标记 ID 来反之亦然，从而利用语言模型中的数值计算。哈希表将很快在 TensorFlow Lite 中启用。它通过在 TensorFlow Lite 格式中本机处理资产文件并提供运算符内核作为 TensorFlow Lite 内置运算符来实现支持。

部署：如何在设备上运行 NLP 模型

使用 TensorFlow Lite 进行推理现在比以前容易得多。您可以使用预构建的推理 API 在 5 行代码内集成您的模型，或者使用实用工具构建您自己的 Android/iOS 推理 API。

使用 TensorFlow Lite 任务库进行简单的模型部署

TensorFlow Lite 任务库 是一款功能强大且易于使用的特定于任务的库，它提供开箱即用的预处理和后处理实用工具，这些实用工具是 ML 推理所必需的，使应用程序开发人员能够轻松地使用 TensorFlow Lite 创建机器学习功能。任务库中支持三个文本 API，它们对应于上述用例和模型

• NLClassifier: 将输入文本分类到一组已知的类别中。
  
• BertNLClassifier: 为 BERT 家族模型优化的文本分类器。
  
• BertQuestionAnswerer: 使用 BERT 家族模型，根据给定段落的內容回答问题。
  

任务库在 Android 和 iOS 平台上均可跨平台使用。以下示例展示了在 Java/Swift 中使用 BertQA 模型进行推理。

// Initialization
BertQuestionAnswerer answerer = BertQuestionAnswerer.createFromFile(androidContext, modelFile);
// Answer a question
List answers = answerer.answer(context, question);

适用于 Android 的 Java 代码

// Initialization
let mobileBertAnswerer =   TFLBertQuestionAnswerer.mobilebertQuestionAnswerer(modelPath: modelPath)
// Answer a question
let answers = mobileBertAnswerer.answer(context: context, question: question)

适用于 iOS 的 Swift 代码

自定义推理 API

如果您的用例不受现有任务库的支持，您也可以利用 Task API 基础设施 并使用相同 仓库 中的常见 NLP 工具（如 Wordpiece 和 Sentencepiece 分词器）构建自己的 C++/Android/iOS 推理 API。

结论

本文介绍了 TensorFlow Lite 中对 NLP 任务的新支持。随着 TensorFlow Lite 的最新更新，开发人员可以轻松地在设备上创建、转换和部署 NLP 模型。我们将继续提供更多有用的工具，并加速从研究到生产的设备上 NLP 模型的开发。我们很乐意听到您的反馈和对新的 NLP 工具和实用程序的建议。请发送电子邮件至 [email protected] 或创建 TensorFlow Lite 支持的 GitHub issue。

致谢

我们感谢 Khanh LeViet、Arun Venkatesan、Max Gubin、Robby Neale、Terry Huang、Peter Young、Gaurav Nemade、Prabhu Kaliamoorthi、Ping Yu、Renjie Liu、Lu Wang、Xunkai Zhang、Yuqi Li、Sijia Ma、Thai Nguyen、Xinying Song、Chung-Ching Chang 和 Shuangfeng Li 对这篇文章的贡献。