利用 TensorFlow Hub 中的预处理模型简化 BERT
文 / 软件工程师 Arno Eigenwillig 和开发技术推广工程师 Luiz GUStavo Martins
BERT 及其他 Transformer 编码器架构在自然语言处理 (NLP) 领域计算矢量空间下的文本表征任务中取得了巨大成功，不仅推进学术领域前沿研究指标的发展，还被广泛应用于 Google 搜索等大型应用。BERT 自最开始便由 TensorFlow 构建，但它最初依赖的并非是 TensorFlow Python 代码来将原始文本转换为模型输入。
今天，我们非常高兴地发布一种更精简的方案，完全基于 TensorFlow 构建的 BERT。由这一解决方案构建的预训练编码器和与之匹配的文本预处理模型现已在 TensorFlow Hub 上提供。现在，只需几行代码，由 TensorFlow 所构建的 BERT 便可以直接处理文本：
# Load BERT and the preprocessing model from TF Hub.
preprocess = hub.load('https://hub.tensorflow.google.cn/tensorflow/bert_en_uncased_preprocess/1')
encoder = hub.load('https://hub.tensorflow.google.cn/tensorflow/bert_en_uncased_L-12_H-768_A-12/3')
# Use BERT on a batch of raw text inputs.
input = preprocess(['Batch of inputs', 'TF Hub makes BERT easy!', 'More text.'])
pooled_output = encoder(input)["pooled_output"]
print(pooled_output)
tf.Tensor(
[[-0.8384154 -0.26902363 -0.3839138 ... -0.3949695 -0.58442086 0.8058556 ]
[-0.8223734 -0.2883956 -0.09359277 ... -0.13833837 -0.6251748 0.88950026]
[-0.9045408 -0.37877116 -0.7714909 ... -0.5112085 -0.70791864 0.92950743]],
shape=(3, 768), dtype=float32)
TensorFlow Hub
https://hub.tensorflow.google.cn/google/collections/bert/1
这些编码器和预处理模型均由 TensorFlow Model Garden 所提供的 NLP 库构建，且以 SavedModel 格式导出至 TensorFlow Hub。在后台，预处理使用 TF.text 库中的 TensorFlow 算子对输入文本执行分词，使您可以构建自己的 TensorFlow 模型，该模型无需在循环中使用 Python 代码即可将原始文本输入转换为预测输出。这不仅加速了计算，消除重复样板代码，减少出错，还实现了对整个文本-输出模型的序列化，使得 BERT 更易于在生产环境中使用。
TensorFlow Model Garden
https://github.com/tensorflow/models/tree/master/official/nlpSavedModel 格式
https://tensorflow.google.cn/hub/tf2_saved_model
初学者教程：处理了一个 情感分析 任务，可在不需要任何特殊自定义的情况下，获得极佳的模型质量。这是使用 BERT 和预处理模型最简单的方法。
https://tensorflow.google.cn/tutorials/text/classify_text_with_bert
进阶教程：处理了来自 GLUE 基准的 NLP 分类任务，并在 TPU 上运行。还展示了如何在多段输入的场景中使用预处理模型。
https://tensorflow.google.cn/tutorials/text/solve_glue_tasks_using_bert_on_tpu
选择 BERT 模型
BERT 模型已使用自监督任务（例如预测上下文句子中的词语），基于大型文本语料库（例如 Wikipedia 文章归档文件）进行预训练。通过此类训练，模型无需标记数据，即可学习文本语义的丰富表征。然而，训练模型需要大量计算：根据 2018 年发表的 BERT 论文，在 16 块 TPU 上需要 4 天。幸运的是，昂贵的预训练一经完成，我们即可在许多不同任务中重复使用这些丰富表征。
8 个 BERT 模型均带有由 BERT 原作者发布的训练权重。
https://hub.tensorflow.google.cn/google/collections/bert/1
24 个 Small BERT 模型采用同一通用架构，但 Transformer 构建块的数量或大小有所减少或缩小，让您能够在速度、规模和质量之间进行权衡。
https://hub.tensorflow.google.cn/google/collections/bert/1
ALBERT：这是四种不同大小的“轻量级 BERT”；通过在层之间共享参数缩减模型大小，但计算时间并未缩短。
https://hub.tensorflow.google.cn/google/collections/albert/1
所有 8 个 BERT Expert 均采用相同的 BERT 架构和大小，可提供不同的预训练域和中间微调任务，以便更接近目标任务。
https://hub.tensorflow.google.cn/google/collections/experts/bert/1
Electra 采用与 BERT 相同的架构（有三种大小可选），但会在类似于生成对抗网络 (GAN) 的配置中作为判别器进行预训练。
https://hub.tensorflow.google.cn/google/collections/electra/1
在具有交谈注意力机制和加门控高斯误差线性单元的 BERT [base、large] 中，我们对 Transformer 架构的核心进行了两项改进。
base
https://hub.tensorflow.google.cn/tensorflow/talkheads_ggelu_bert_en_base/1large
https://hub.tensorflow.google.cn/tensorflow/talkheads_ggelu_bert_en_large/1
Lambert 已使用 LAMB 优化器以及来自 RoBERTa 的若干技术进行训练。
https://hub.tensorflow.google.cn/tensorflow/lambert_en_uncased_L-24_H-1024_A-16/1
逐步添加更多模型，敬请期待……
这些模型均为 BERT 编码器。通过上方链接将转到 TF Hub 上的相应文档，文档中指出了适用于每种 BERT 编码器的预处理模型。
我们建议开发者访问这些模型页面，以详细了解每种模型所针对的不同应用。编码器均使用通用接口，因此您可以通过更改编码器模型网址及其预处理，基于特定任务实验并比较不同编码器的性能。
预处理模型
每种 BERT 编码器都有一个与之匹配的预处理模型。该模型使用 TF.text 库提供的 TensorFlow 算子，将原始文本转换为编码器需要的数字输入张量。与使用纯 Python 进行预处理不同，这些算子可以成为 TensorFlow 模型的一部分，能够直接处理输入文本。TF Hub 中的所有预处理模型均已配置词汇表及其关联的文本标准化逻辑，无需进一步设置。
上文介绍了使用预处理模型最简单的方法。接下来我们将进一步介绍：
preprocess = hub.load('https://hub.tensorflow.google.cn/tensorflow/bert_en_uncased_preprocess/1')
input = preprocess(["This is an amazing movie!"])
{'input_word_ids': <tf.Tensor: shape=(1, 128), dtype=int32, numpy=
array([[ 101, 2023, 2003, 2019, 6429, 3185, 999, 102, 0, ...]])>,
'input_mask': <tf.Tensor: shape=(1, 128), dtype=int32, numpy=
array([[ 1, 1, 1, 1, 1, 1, 1, 1, 0, ...,]])>,
'input_type_ids': <tf.Tensor: shape=(1, 128), dtype=int32, numpy=
array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, ...,]])>}
如此调用 preprocess() 会将原始文本输入转换为用于 BERT 编码器的固定长度输入序列。可以看到，该序列由 input_word_ids张量与每个输入分词后的数字 id 组成，其中包括起始标记、结束标记和填充标记；以及两个辅助张量：input_mask （用于区分填充标记和非填充标记）和每个标记的 input_type_ids （用于区分每个输入的多个文本段，这一部分我们稍后再讨论）。
同一个预处理 SavedModel 还将再提供一个更为精细的 API，支持将一个或两个不同的文本段放入编码器的同一输入序列中。我们来看一个句子蕴涵 (sentence entailment) 任务，其中 BERT 用于预测一个前提是否蕴涵某个假设：
text_premises = ["The fox jumped over the lazy dog.",
"Good day."]
tokenized_premises = preprocess.tokenize(text_premises)
<tf.RaggedTensor
[[[1996], [4419], [5598], [2058], [1996], [13971], [3899], [1012]],
[[2204], [2154], [1012]]]>
text_hypotheses = ["The dog was lazy.", # Entailed.
"Axe handle!"] # Not entailed.
tokenized_hypotheses = preprocess.tokenize(text_hypotheses)
<tf.RaggedTensor
[[[1996], [3899], [2001], [13971], [1012]],
[[12946], [5047], [999]]]>
每次分词都会生成一个数字标记 ID 的不规则张量 (RaggedTensor)，完整表示每个文本输入。如果某些前提和假设的组合过长，超过了下一步中 BERT 输入的 seq_length， 您可以再次进行预处理，比如修剪文本段或将文本段分割为多个编码器输入。
RaggedTensor
https://tensorflow.google.cn/guide/ragged_tensor
然后，分词的输入就会被打包成一个用于 BERT 编码器的固定长度输入序列：
encoder_inputs = preprocess.bert_pack_inputs(
[tokenized_premises, tokenized_hypotheses],
seq_length=18) # Optional argument, defaults to 128.
{'input_word_ids': <tf.Tensor: shape=(2, 18), dtype=int32, numpy=
array([[ 101, 1996, 4419, 5598, 2058, 1996, 13971, 3899, 1012,
102, 1996, 3899, 2001, 13971, 1012, 102, 0, 0],
[ 101, 2204, 2154, 1012, 102, 12946, 5047, 999, 102,
0, 0, 0, 0, 0, 0, 0, 0, 0]])>,
'input_mask': <tf.Tensor: shape=(2, 18), dtype=int32, numpy=
array([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]])>,
'input_type_ids': <tf.Tensor: shape=(2, 18), dtype=int32, numpy=
array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0],
[0, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]])>}
打包将生成常见的 input_word_ids、 input_mask 和 input_type_ids （第一个和第二个输入分别是 0 和 1）字典。所有输出都有一个统一的 seq_length （默认情况下为 128 个字符）。超过 seq_length 的输入将会在打包过程中被截断成大小相同的若干段。
加速模型训练
TensorFlow Hub 提供单独的 BERT 编码器和预处理模型以加速训练，尤其是 TPU 上的训练。
张量处理单元 (TPU) 是 Google 开发的自定义加速器硬件，在大规模机器学习计算（如微调 BERT 所需要的计算）中的表现卓越。TPU 对密集张量进行操作，并期望主机 CPU 已将字符串等可变长度数据转换为固定大小的张量。
通过分离 BERT 编码器模型和与之关联的预处理模型，可将编码器微调计算分发到 TPU，作为模型训练的一部分，同时在主机 CPU 上执行预处理模型。预处理计算可使用 tf.data.Dataset.map() 在数据集上异步运行，同时准备密集输出以供 TPU 上的编码器模型使用。此类异步预处理也可提高其他加速器的性能。
我们的进阶 BERT 教程可以在 Colab 中运行，运行时使用 TPU Worker 并对其进行端到端演示。
进阶 BERT 教程
https://tensorflow.google.cn/tutorials/text/solve_glue_tasks_using_bert_on_tpu
总结
在 TensorFlow 中使用 BERT 及类似模型变得更加简单。TensorFlow Hub 提供了丰富多样的预训练 BERT 编码器和文本预处理模型，它们只需几行代码即可轻松使用。
丰富多样
https://hub.tensorflow.google.cn/google/collections/transformer_encoders_text/1
查看我们的交互式初学者和进阶教程，详细了解如何使用模型对句子和句对进行分类。欢迎投稿与我们分享您使用全新 BERT 模型构建了哪些新奇应用。
初学者
https://tensorflow.google.cn/tutorials/text/classify_text_with_bert进阶教程
https://tensorflow.google.cn/tutorials/text/solve_glue_tasks_using_bert_on_tpu投稿
http://services.google.cn/fb/forms/TFCS/?channel=wechat
致谢
我们要感谢许多同事为本工作所做的贡献。
全新预处理模型由 Chen Chen、Terry Huang、Mark Omernick 和 Rajagopal Ananthanarayanan 合作创建。
值此之际，更多 BERT 模型也已由 Sebastian Ebert (Small BERT)、Le Hou 和 Hongkun Yu（Lambert、Talking Heads）发布到 TF Hub。
Mark Daoust、Josh Gordon 和 Elizabeth Kemp 极大地改善了本文及相关教程中材料的呈现方式。
了解更多请点击 “阅读原文” 访问官网。