TensorFlow Models NLP：models 预置模型体系——BertClassifier 到 T5Transformer 的七种可训练 Keras 模型深度解析

在 official/nlp/modeling/models/README.md 中，官方 NLP 建模库将 “Model（模型）” 定义为“由 tf.keras 层与模型组合而成、可直接训练的对象”，并提供了多个预置（canned）模型用于训练 encoder 网络。这些模型既是方便用户快速搭建任务的便捷函数，也是官方认可的“规范示例（canonical examples）”。本篇将逐一剖析这些预置模型的网络结构、构造参数与训练/推理行为，并结合仓库中的源码与测试用例说明其实现细节，帮助你在分类、标注、问答、预训练、检索、序列到序列生成等典型 NLP 场景中正确选用并组装这些模型。

一、models 模块的整体定位

从包定义 official/nlp/modeling/models/init.py 可以看到，该模块对外暴露的模型族包括：

• 句子级分类：BertClassifier（bert_classifier.py）
• 词元级分类：BertTokenClassifier（bert_token_classifier.py）
• 区间（span）标注：BertSpanLabeler（bert_span_labeler.py）
• BERT 预训练：BertPretrainer / BertPretrainerV2（bert_pretrainer.py）
• ELECTRA 预训练：ElectraPretrainer（electra_pretrainer.py）
• 双塔检索：DualEncoder（dual_encoder.py）
• 原始 Transformer 序列到序列：Seq2SeqTransformer（seq2seq_transformer.py）
• T5：T5Transformer（t5.py）
• XLNet 系列：XLNetClassifier、XLNetPretrainer、XLNetSpanLabeler（xlnet.py）

需要区分两个层级：network（网络） 是可复用的 Transformer 编码器栈（如 networks/bert_encoder.py 中的 BertEncoder），model（模型） 则是在某个 network 之上叠加任务头（classification head、token head、span head、masked LM 头）后形成的完整 tf_keras.Model。下文的每个模型都遵循这一“network + 任务头”的组合模式。

二、BertClassifier：基于 CLS 池化输出的句子分类/回归模型

BertClassifier 实现了围绕 Transformer 编码器的经典 BERT 分类网络结构（对应论文 "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding"）。它的用途是：传入一个 transformer 网络和一个类别数，即得到一个可直接 fit 的分类（或回归）模型。

2.1 构造参数

参数	默认值	说明
network	必填	transformer 网络，需输出 sequence output 与 classification output，并暴露 get_embedding_table 方法
num_classes	必填	分类头输出的类别数；设为 1 时该模型即用作回归模型
initializer	'glorot_uniform'	分类网络的权重初始化器
dropout_rate	0.1	cls 头的 dropout 概率
use_encoder_pooler	True	是否使用编码器内置的 pooler 层（即 CLS 池化输出）
head_name	'sentence_prediction'	分类头命名
cls_head	None	可选的自定义分类头层实例；一旦设置，num_classes、initializer、dropout_rate、use_encoder_pooler、head_name 均被忽略

2.2 关键实现行为

从 BertClassifier.init 的源码可以看到：

1. Functional API 组装：模型通过保存 network.inputs 句柄，用 network 自身的输入张量调用 network 得到输出，再在 __init__ 末尾调用 super().__init__(inputs=..., outputs=predictions)，使对象具备完整 Functional API 模型的属性（源码注释中引用了 b/164516224 说明该构造顺序的必要性）。
2. 两种输入来源：use_encoder_pooler=True 时取 outputs[1]（或字典形式的 pooled_output）并先过一层 Dropout；False 时取序列输出 outputs[0]（或 sequence_output）。
3. 分类头：默认使用 layers/cls_head.py 中的 ClassificationHead，注意 inner_dim=0 if use_encoder_pooler else cls_inputs.shape[-1]——使用编码器 pooler 时头内不再额外加隐藏维；未使用 pooler 时则以序列宽度作为 inner_dim。
4. 配置序列化：config 以 collections.namedtuple 存储（而非 dict），源码注释解释其动机是“TF 不会跟踪不含 Trackable 的不可变属性”，从而保持与旧版本检查点的兼容性。模型同时实现 get_config / from_config，并被 register_keras_serializable(package='Text') 注册。
5. 检查点映射：checkpoint_items 属性将 encoder 指向 network、并将分类头内部可训练项以 头名.键名 拼接后返回，便于与 BERT 预训练检查点对齐加载。

2.3 测试用例给出的用法验证

bert_classifier_test.py 展示了标准装配流程：

test_network = networks.BertEncoder(vocab_size=100, num_layers=2, dict_outputs=dict_outputs)
bert_trainer_model = bert_classifier.BertClassifier(test_network, num_classes=num_classes)

word_ids = tf_keras.Input(shape=(sequence_length,), dtype=tf.int32)
mask = tf_keras.Input(shape=(sequence_length,), dtype=tf.int32)
type_ids = tf_keras.Input(shape=(sequence_length,), dtype=tf.int32)

cls_outs = bert_trainer_model([word_ids, mask, type_ids])
# 输出形状校验为 [None, num_classes]

测试还验证了两个易被忽视的能力：

• 传入 cls_head=layers.GaussianProcessClassificationHead(inner_dim=0, num_classes=num_classes) 时模型可正常前向（即 SNNGP 式自定义头与自定义 num_classes 组合工作）；
• get_config() 与 BertClassifier.from_config(config) 往返一致，且 config 可强制转为 JSON（to_json()），证明模型可完整序列化。

三、BertTokenClassifier：在序列输出上做词元级分类

BertTokenClassifier 在 sequence output 上叠加一个单层 Dense 分类头，适用于 NER、POS 等词元级任务。构造参数与 BertClassifier 类似，但多一个输出风格开关：

参数	默认值	说明
network	必填	同上，需输出 sequence output 与 classification output
num_classes	必填	每个词元位置的类别数
initializer	'glorot_uniform'	Dense 头的初始化器
output	'logits'	输出风格：logits（原始 logits）或 predictions（经 log_softmax 的预测）
dropout_rate	0.1	分类头前对 sequence output 的 dropout
output_encoder_outputs	False	是否在最终输出字典中额外附带 encoder_outputs（编码器的序列输出）

从 源码 可见其结构：对 sequence_output 施加 Dropout 后送入命名为 predictions/transform/logits 的 Dense 层；output='logits' 时输出字典仅含 logits 键，output='predictions' 时键名变为 predictions 且数值为 tf.nn.log_softmax 结果，其他取值直接抛出 ValueError。若设置 output_encoder_outputs=True，输出字典还会附加 encoder_outputs 键，方便下游（如解码、可视化）拿到词元表示。模型同样实现 checkpoint_items（仅映射 encoder）、get_config、from_config，可直接 tf.keras.models.save_model 保存为 SavedModel。

四、BertSpanLabeler：SQuAD 式起点-终点区间预测

BertSpanLabeler 实现“单一区间起点-终点预测器”：对每个词元位置输出两个值——start token 的 logit 与 end token 的 logit，适合 SQuAD 风格的可抽取式问答。

构造参数只有三个：network、initializer='glorot_uniform'、output（'logits' 或 'predictions'）。

实现上有两个值得注意的细节（见 源码）：

1. 多层编码器输出的兼容：sequence_output 可能是 list（当编码器被配置为返回所有层输出时），此时取 sequence_output[-1] 作为最后一层表示，这使该模型可同时适配普通编码器与 return-all-outputs 的编码器；
2. 命名输出技巧：start/end logits 分别经过 tf_keras.layers.Lambda(tf.identity, name='start_positions') 和 name='end_positions' 包装。源码注释明确说明目的：“通过显式命名输出张量，可以在 Keras 的 fit/predict/evaluate 调用中使用字符串键字典”。也就是说，训练/评估时你可以用 {'start_positions': ..., 'end_positions': ...} 这样的字典作为 y 或 loss 的键。

span 头本体来自 networks/span_labeling.py 中的 SpanLabeling 网络，其宽度取自 sequence_output.shape[-1]。

五、BertPretrainer / BertPretrainerV2：预训练目标（Masked LM + NSP）

bet_pretrainer.py 提供两个预训练模型，二者都“在 Transformer 编码器之上实例化训练目标所需的网络”。

5.1 BertPretrainer（经典 BERT 预训练）

构造参数（见 源码）：

参数	说明
network	transformer 网络，需输出 sequence output 与 classification output
num_classes	句对分类（NSP）网络的类别数
num_token_predictions	Masked LM 头预测的 token 数
embedding_table	若为 None，则调用 network.get_embedding_table()
activation	Masked LM 网络激活函数；None 表示不使用
initializer	'glorot_uniform'，经 tf_utils.clone_initializer 克隆后分别传入两个头
output	'logits' 或 'predictions'

关键行为：

• 额外输入：模型输入 = 编码器全部输入 + 一个新构造的 masked_lm_positions（shape=(num_token_predictions,) 的 tf.int32 Input）。若静态可知的序列长度小于 num_token_predictions 会直接抛出 ValueError，这是构造时的一个显式合法性检查；
• Masked LM 头：使用 layers/masked_lm.py 的 MaskedLM 层，name='cls/predictions'，与 BERT 原模型变量命名对齐；
• 句对分类头：使用 networks/classification.py 的 Classification 网络，name='classification'，作用在 CLS 池化输出上；
• 输出字典：dict(masked_lm=lm_outputs, classification=sentence_outputs)，即训练时损失可分别绑定到两个键。

5.2 BertPretrainerV2（推荐版本）

BertPretrainer 的 docstring 明确提示：“Please use the new BertPretrainerV2 for your projects.”。BertPretrainerV2 带有 @gin.configurable 装饰器（可通过 gin 配置注入），参数如下：

参数	说明
encoder_network	编码器网络（构造时会主动调用一次以强制 build 权重）
mlm_activation / mlm_initializer	Masked LM 的激活与初始化器
classification_heads	可选的额外分类头列表（要求各头 name 唯一，否则抛 ValueError），例如加入一个 NSP 头
customized_masked_lm	自定义 Masked LM 层；若提供则忽略 mlm_activation 与 mlm_initializer
name	模型名，默认 'bert'

与 V1 的差异点：

• masked_lm_positions 的 shape 放宽为 (None,)（int32），并支持以字典形式并入编码器输入（inputs['masked_lm_positions'] = ...）；
• call() 在推理场景下允许缺失 masked_lm_positions，此时跳过 MLM 前向（源码注释：“Inference may not have masked_lm_positions and mlm_logits is not needed.”）；
• 输出字典包含 pooled_output、sequence_output、可选 encoder_outputs（多层输出时）、mlm_logits，以及各分类头以其 name 为键的输出，结构比 V1 更适合多任务预训练。

六、DualEncoder：面向检索的双塔编码器

DualEncoder 依据 "Language-agnostic BERT Sentence Embedding"（LaBSE）的双塔结构实现：同一个 transformer 网络对左右两条序列分别编码，比较各自的句向量（pooled output），适用于句检索/相似度训练。

构造参数（见 源码）：

参数	默认值	说明
network	必填	输出 encoding 输出的 transformer 网络
max_seq_length	32	transformer 允许的最大序列长度
normalize	True	是否对 pooled 输出做 L2 归一化（tf.nn.l2_normalize）
logit_scale	1.0	训练时对点积的缩放因子
logit_margin	0.0	训练时正负样本之间的 margin
output	'logits'	'logits'（双输入、输出 left/right logits）或 'predictions'（单输入、直接输出嵌入）

实现要点：

• 输入命名按用途切换：output='logits'（训练）时左塔输入名为 left_word_ids/left_mask/left_type_ids、右塔为 right_*；output='predictions'（推理/嵌入导出）时输入名改为 input_word_ids/input_mask/input_type_ids——源码注释说明这是为了与旧版 BERT Hub 模块的输入名保持一致；
• 点积层：使用 layers.MatMulWithMargin（name='dot_product'）同时产出 left_logits 与 right_logits，二者结合 logit_scale 与 logit_margin 实现对称的对比学习目标；
• 嵌入输出：output='predictions' 时输出为 dict(sequence_output=..., pooled_output=left_encoded)，同样出于与旧 BERT Hub 模块输出名一致的目的。

该模型在 official/projects/labse 项目中有完整实验配置，可结合查看双塔检索的端到端训练方式。

七、Seq2SeqTransformer：原始 Transformer 机器翻译模型

Seq2SeqTransformer 依据 "Attention Is All You Need" 论文（arXiv:1706.03762）实现，是官方库中序列到序列任务的参考实现。它由三部分组成：

7.1 顶层模型 Seq2SeqTransformer

构造参数（见 源码）：

参数	默认值	说明
vocab_size	33708	词表大小
embedding_width	512	嵌入/隐层宽度
dropout_rate	0.0	dropout 概率
padded_decode	False	是否使用按 decode_max_length 填充的解码（TPU 场景）
decode_max_length	None	解码最大步数；padded_decode=False 时若未指定则取 源长度 + extra_decode_length
extra_decode_length	0	束搜索额外运行的步数
beam_size	4	束搜索的束宽
alpha	0.6	束搜索长度归一化强度
encoder_layer / decoder_layer	None	需外部传入已初始化的编码器/解码器层实例
eos_id	1（EOS_ID）	句末 token id

内部构建：

• 词嵌入用 layers.OnDeviceEmbedding，initializer 为标准差 embedding_width**-0.5 的正态分布，scale_factor = embedding_width**0.5（即原论文中嵌入向量乘以 √d 的做法）；
• 位置信息用 layers.RelativePositionEmbedding（相对位置编码）叠加到嵌入上；
• 输出投影通过 _embedding_linear 以词嵌入矩阵的转置作为线性变换权重，实现论文中的 weight tying。

call() 的双模式行为：

• 训练（提供 targets）：将 targets 右移一位并截去末位作为 decoder 输入，构造下三角自注意力掩码（tf.linalg.band_part），返回 (batch_size, target_length, vocab_size) 的 float32 logits（源码显式 tf.cast(logits, tf.float32) 以避免混合精度下的数值问题）；
• 推理（targets 为 None）：构造每层 key/value 的 cache（形状 [batch, decode_len, heads, width//heads]），把编码器输出与 encoder-decoder 注意力掩码也放入 cache，然后调用 beam_search.sequence_beam_search，返回 {'outputs': (batch, decoded_len), 'scores': (batch, 1)}（取束中第 0 条即最高分序列）。

输入字典要求 inputs 与 embedded_inputs+input_masks 二选一（见 _parse_inputs），padding 位置以 token id 0 判定（boolean_mask = tf.not_equal(sources, 0)）。

7.2 TransformerEncoder / TransformerDecoder

两个堆叠层（源码）参数一致：num_layers=6、num_attention_heads=8、intermediate_size=2048、activation='relu'、dropout_rate=0.0、attention_dropout_rate=0.0、use_bias=False、norm_first=True、norm_epsilon=1e-6、intermediate_dropout=0.0，与论文超参对齐。

• 编码器 build() 时按层实例化 layers.TransformerEncoderBlock，注意力初始化器用 attention_initializer（Glorot uniform，limit = sqrt(6/(2*hidden_size))），最后过一层 LayerNormalization；
• 解码器每层为 layers.TransformerDecoderBlock，支持通过 self_attention_cls / cross_attention_cls 按层注入自定义注意力类（类或函数均可），并支持 cache 参数走快速解码路径；return_all_decoder_outputs=True 可返回逐层归一化输出（源码注释指出这便于引入逐层辅助损失）。

八、T5Transformer：与官方 T5 架构/checkpoint 兼容的实现

t5.py 实现了独立的 T5 模型，面向 seq-to-seq 任务。文件头注释声明了两点关键事实：

1. 公开接口只有两个：T5TransformerParams（t5.py#L1006，dataclass 形式的参数集）与 T5Transformer（t5.py#L1379），其余模块（Module、Embed、make_attention_mask、make_causal_mask 等）属于实现细节，不建议下游库直接依赖；
2. checkpoint 兼容性：模型与已发布的 T5 架构及转换后的 checkpoint 兼容。

实现风格与前面几个模型不同：T5 各模块以 tf.Module 实现（而非 tf.keras）。因此 README 特别给出使用指引——若要放进 Keras 训练流程，应在自定义 Keras 层的 __init__ 中实例化 T5 模块、在 call 中调用它们，用 Keras 层做一层包装。文件内还提供了一些有代表性的内部机制，例如：

• Embed 模块默认 one_hot=True：用 tf.one_hot 与嵌入矩阵的矩阵乘法代替 embedding_lookup，以获得更稠密的梯度路径；one_hot=False 时则用 embedding_lookup 配合 dense_gradient（一个 tf.custom_gradient 操作，将 IndexedSlices 梯度转为稠密张量）；
• make_causal_mask 基于下标广播生成 [batch..., 1, len, len] 的因果掩码，供 decoder 自注意力使用。

九、通用工程特性：序列化、检查点与测试

跨这些模型，仓库保持了几项一致的可工程化约定，选型时可作为可靠预期：

• Keras 可序列化：BERT 系模型均带 @tf_keras.utils.register_keras_serializable(package='Text')，实现 get_config / from_config，可 to_json / from_json 往返（bert_classifier_test.py#L84-L103 对 config 往返与 JSON 化做了断言）；
• 检查点项映射：通过 checkpoint_items 属性把 encoder 等子模块映射为可保存/可加载条目，使微调模型能加载预训练编码器权重而不受外层模型结构差异影响；
• 输出风格统一：涉及预测头的模型普遍支持 output='logits' | 'predictions'，predictions 模式对分类头做 log_softmax 变换，便于与交叉熵损失直接对接；
• 命名对齐原论文：关键层命名（cls/predictions、predictions/transform/logits、dot_product 等）刻意与 BERT/T5/双塔论文及社区 checkpoint 的变量名保持一致，为检查点互操作留下空间。

每个模型都配有同名 _test.py（如 bert_pretrainer_test.py、dual_encoder_test.py、t5_test.py），覆盖构建、前向、序列化等契约，可作为集成时最贴近预期的行为参照。

十、如何选型与组合

结合上述源码，可以把七个模型映射到常见任务：

任务	模型	输入/输出要点
文本分类/回归	BertClassifier	[word_ids, mask, type_ids] → [batch, num_classes]；num_classes=1 即回归
NER/词元分类	BertTokenClassifier	→ 词元位置 logits（可要求附带 encoder_outputs）
可抽取式问答	BertSpanLabeler	→ start_positions / end_positions 两路 logit
BERT 预训练	BertPretrainer（或 V2）	额外输入 masked_lm_positions → {masked_lm, classification}
句向量检索	DualEncoder	双输入 → {left_logits, right_logits}；嵌入模式输出 pooled_output
机器翻译/通用 seq2seq	Seq2SeqTransformer	训练返回 targets 的 logits；推理走束搜索返回 {outputs, scores}
T5 风格 seq2seq	T5Transformer	tf.Module 实现，需自行包装进 Keras 层

使用上的通用模式是：先实例化一个 network（如 networks.BertEncoder(vocab_size=..., num_layers=...)），再把它交给对应的 model 构造器——测试代码中的 “BertEncoder + BertClassifier”组合 就是最小可复现范例。理解这套“network 定义骨干、model 定义任务头、checkpoint_items 对齐权重、config 支持序列化”的分层约定后，你可以按同样方式把其他编码器（ALBERT、MobileBERT 等，见 official/nlp/modeling/networks）与这些任务头自由组合，构建出新的预置模型。