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

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

【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models

在 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的源码可以看到：

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

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

多层编码器输出的兼容：sequence_output可能是 list（当编码器被配置为返回所有层输出时），此时取sequence_output[-1]作为最后一层表示，这使该模型可同时适配普通编码器与 return-all-outputs 的编码器；
命名输出技巧：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.int32Input）。若静态可知的序列长度小于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 newBertPretrainerV2for 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 任务。文件头注释声明了两点关键事实：

公开接口只有两个：T5TransformerParams（t5.py#L1006，dataclass 形式的参数集）与T5Transformer（t5.py#L1379），其余模块（Module、Embed、make_attention_mask、make_causal_mask等）属于实现细节，不建议下游库直接依赖；
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）与这些任务头自由组合，构建出新的预置模型。

【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析