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

TensorFlow Models NLPmodels 预置模型体系——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 可以看到该模块对外暴露的模型族包括句子级分类BertClassifierbert_classifier.py词元级分类BertTokenClassifierbert_token_classifier.py区间span标注BertSpanLabelerbert_span_labeler.pyBERT 预训练BertPretrainer/BertPretrainerV2bert_pretrainer.pyELECTRA 预训练ElectraPretrainerelectra_pretrainer.py双塔检索DualEncoderdual_encoder.py原始 Transformer 序列到序列Seq2SeqTransformerseq2seq_transformer.pyT5T5Transformert5.pyXLNet 系列XLNetClassifier、XLNetPretrainer、XLNetSpanLabelerxlnet.py需要区分两个层级network网络是可复用的 Transformer 编码器栈如 networks/bert_encoder.py 中的BertEncodermodel模型则是在某个 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 时该模型即用作回归模型initializerglorot_uniform分类网络的权重初始化器dropout_rate0.1cls 头的 dropout 概率use_encoder_poolerTrue是否使用编码器内置的 pooler 层即 CLS 池化输出head_namesentence_prediction分类头命名cls_headNone可选的自定义分类头层实例一旦设置num_classes、initializer、dropout_rate、use_encoder_pooler、head_name均被忽略2.2 关键实现行为从 BertClassifier.init的源码可以看到Functional API 组装模型通过保存network.inputs句柄用 network 自身的输入张量调用 network 得到输出再在__init__末尾调用super().__init__(inputs..., outputspredictions)使对象具备完整 Functional API 模型的属性源码注释中引用了 b/164516224 说明该构造顺序的必要性。两种输入来源use_encoder_poolerTrue时取outputs[1]或字典形式的pooled_output并先过一层 DropoutFalse时取序列输出outputs[0]或sequence_output。分类头默认使用 layers/cls_head.py 中的ClassificationHead注意inner_dim0 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(packageText)注册。检查点映射checkpoint_items属性将encoder指向 network、并将分类头内部可训练项以头名.键名拼接后返回便于与 BERT 预训练检查点对齐加载。2.3 测试用例给出的用法验证bert_classifier_test.py 展示了标准装配流程test_network networks.BertEncoder(vocab_size100, num_layers2, dict_outputsdict_outputs) bert_trainer_model bert_classifier.BertClassifier(test_network, num_classesnum_classes) word_ids tf_keras.Input(shape(sequence_length,), dtypetf.int32) mask tf_keras.Input(shape(sequence_length,), dtypetf.int32) type_ids tf_keras.Input(shape(sequence_length,), dtypetf.int32) cls_outs bert_trainer_model([word_ids, mask, type_ids]) # 输出形状校验为 [None, num_classes]测试还验证了两个易被忽视的能力传入cls_headlayers.GaussianProcessClassificationHead(inner_dim0, num_classesnum_classes)时模型可正常前向即 SNNGP 式自定义头与自定义num_classes组合工作get_config()与BertClassifier.from_config(config)往返一致且 config 可强制转为 JSONto_json()证明模型可完整序列化。三、BertTokenClassifier在序列输出上做词元级分类BertTokenClassifier 在 sequence output 上叠加一个单层 Dense 分类头适用于 NER、POS 等词元级任务。构造参数与BertClassifier类似但多一个输出风格开关参数默认值说明network必填同上需输出 sequence output 与 classification outputnum_classes必填每个词元位置的类别数initializerglorot_uniformDense 头的初始化器outputlogits输出风格logits原始 logits或predictions经log_softmax的预测dropout_rate0.1分类头前对 sequence output 的 dropoutoutput_encoder_outputsFalse是否在最终输出字典中额外附带encoder_outputs编码器的序列输出从 源码 可见其结构对sequence_output施加 Dropout 后送入命名为predictions/transform/logits的 Dense 层outputlogits时输出字典仅含logits键outputpredictions时键名变为predictions且数值为tf.nn.log_softmax结果其他取值直接抛出ValueError。若设置output_encoder_outputsTrue输出字典还会附加encoder_outputs键方便下游如解码、可视化拿到词元表示。模型同样实现checkpoint_items仅映射encoder、get_config、from_config可直接tf.keras.models.save_model保存为 SavedModel。四、BertSpanLabelerSQuAD 式起点-终点区间预测BertSpanLabeler 实现“单一区间起点-终点预测器”对每个词元位置输出两个值——start token 的 logit 与 end token 的 logit适合 SQuAD 风格的可抽取式问答。构造参数只有三个network、initializerglorot_uniform、outputlogits或predictions。实现上有两个值得注意的细节见 源码多层编码器输出的兼容sequence_output可能是 list当编码器被配置为返回所有层输出时此时取sequence_output[-1]作为最后一层表示这使该模型可同时适配普通编码器与 return-all-outputs 的编码器命名输出技巧start/end logits 分别经过tf_keras.layers.Lambda(tf.identity, namestart_positions)和nameend_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 NSPbet_pretrainer.py 提供两个预训练模型二者都“在 Transformer 编码器之上实例化训练目标所需的网络”。5.1 BertPretrainer经典 BERT 预训练构造参数见 源码参数说明networktransformer 网络需输出 sequence output 与 classification outputnum_classes句对分类NSP网络的类别数num_token_predictionsMasked LM 头预测的 token 数embedding_table若为 None则调用network.get_embedding_table()activationMasked LM 网络激活函数None 表示不使用initializerglorot_uniform经tf_utils.clone_initializer克隆后分别传入两个头outputlogits或predictions关键行为额外输入模型输入 编码器全部输入 一个新构造的masked_lm_positionsshape(num_token_predictions,)的tf.int32Input。若静态可知的序列长度小于num_token_predictions会直接抛出ValueError这是构造时的一个显式合法性检查Masked LM 头使用 layers/masked_lm.py 的MaskedLM层namecls/predictions与 BERT 原模型变量命名对齐句对分类头使用 networks/classification.py 的Classification网络nameclassification作用在 CLS 池化输出上输出字典dict(masked_lmlm_outputs, classificationsentence_outputs)即训练时损失可分别绑定到两个键。5.2 BertPretrainerV2推荐版本BertPretrainer的 docstring 明确提示“Please use the newBertPretrainerV2for your projects.”。BertPretrainerV2 带有gin.configurable装饰器可通过 gin 配置注入参数如下参数说明encoder_network编码器网络构造时会主动调用一次以强制 build 权重mlm_activation/mlm_initializerMasked LM 的激活与初始化器classification_heads可选的额外分类头列表要求各头name唯一否则抛ValueError例如加入一个 NSP 头customized_masked_lm自定义 Masked LM 层若提供则忽略mlm_activation与mlm_initializername模型名默认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 EmbeddingLaBSE的双塔结构实现同一个 transformer 网络对左右两条序列分别编码比较各自的句向量pooled output适用于句检索/相似度训练。构造参数见 源码参数默认值说明network必填输出 encoding 输出的 transformer 网络max_seq_length32transformer 允许的最大序列长度normalizeTrue是否对 pooled 输出做 L2 归一化tf.nn.l2_normalizelogit_scale1.0训练时对点积的缩放因子logit_margin0.0训练时正负样本之间的 marginoutputlogitslogits双输入、输出 left/right logits或predictions单输入、直接输出嵌入实现要点输入命名按用途切换outputlogits训练时左塔输入名为left_word_ids/left_mask/left_type_ids、右塔为right_*outputpredictions推理/嵌入导出时输入名改为input_word_ids/input_mask/input_type_ids——源码注释说明这是为了与旧版 BERT Hub 模块的输入名保持一致点积层使用layers.MatMulWithMarginnamedot_product同时产出left_logits与right_logits二者结合logit_scale与logit_margin实现对称的对比学习目标嵌入输出outputpredictions时输出为dict(sequence_output..., pooled_outputleft_encoded)同样出于与旧 BERT Hub 模块输出名一致的目的。该模型在 official/projects/labse 项目中有完整实验配置可结合查看双塔检索的端到端训练方式。七、Seq2SeqTransformer原始 Transformer 机器翻译模型Seq2SeqTransformer 依据 Attention Is All You Need 论文arXiv:1706.03762实现是官方库中序列到序列任务的参考实现。它由三部分组成7.1 顶层模型 Seq2SeqTransformer构造参数见 源码参数默认值说明vocab_size33708词表大小embedding_width512嵌入/隐层宽度dropout_rate0.0dropout 概率padded_decodeFalse是否使用按decode_max_length填充的解码TPU 场景decode_max_lengthNone解码最大步数padded_decodeFalse时若未指定则取源长度 extra_decode_lengthextra_decode_length0束搜索额外运行的步数beam_size4束搜索的束宽alpha0.6束搜索长度归一化强度encoder_layer/decoder_layerNone需外部传入已初始化的编码器/解码器层实例eos_id1EOS_ID句末 token id内部构建词嵌入用layers.OnDeviceEmbeddinginitializer为标准差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_inputsinput_masks二选一见_parse_inputspadding 位置以 token id 0 判定boolean_mask tf.not_equal(sources, 0)。7.2 TransformerEncoder / TransformerDecoder两个堆叠层源码参数一致num_layers6、num_attention_heads8、intermediate_size2048、activationrelu、dropout_rate0.0、attention_dropout_rate0.0、use_biasFalse、norm_firstTrue、norm_epsilon1e-6、intermediate_dropout0.0与论文超参对齐。编码器build()时按层实例化layers.TransformerEncoderBlock注意力初始化器用attention_initializerGlorot uniformlimit sqrt(6/(2*hidden_size))最后过一层LayerNormalization解码器每层为layers.TransformerDecoderBlock支持通过self_attention_cls/cross_attention_cls按层注入自定义注意力类类或函数均可并支持cache参数走快速解码路径return_all_decoder_outputsTrue可返回逐层归一化输出源码注释指出这便于引入逐层辅助损失。八、T5Transformer与官方 T5 架构/checkpoint 兼容的实现t5.py 实现了独立的 T5 模型面向 seq-to-seq 任务。文件头注释声明了两点关键事实公开接口只有两个T5TransformerParamst5.py#L1006dataclass 形式的参数集与T5Transformert5.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_hotTrue用tf.one_hot与嵌入矩阵的矩阵乘法代替embedding_lookup以获得更稠密的梯度路径one_hotFalse时则用embedding_lookup配合dense_gradient一个tf.custom_gradient操作将IndexedSlices梯度转为稠密张量make_causal_mask基于下标广播生成[batch..., 1, len, len]的因果掩码供 decoder 自注意力使用。九、通用工程特性序列化、检查点与测试跨这些模型仓库保持了几项一致的可工程化约定选型时可作为可靠预期Keras 可序列化BERT 系模型均带tf_keras.utils.register_keras_serializable(packageText)实现get_config/from_config可to_json/from_json往返bert_classifier_test.py#L84-L103 对 config 往返与 JSON 化做了断言检查点项映射通过checkpoint_items属性把encoder等子模块映射为可保存/可加载条目使微调模型能加载预训练编码器权重而不受外层模型结构差异影响输出风格统一涉及预测头的模型普遍支持outputlogits | predictionspredictions模式对分类头做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_classes1即回归NER/词元分类BertTokenClassifier→ 词元位置 logits可要求附带 encoder_outputs可抽取式问答BertSpanLabeler→start_positions/end_positions两路 logitBERT 预训练BertPretrainer或 V2额外输入masked_lm_positions→{masked_lm, classification}句向量检索DualEncoder双输入 →{left_logits, right_logits}嵌入模式输出pooled_output机器翻译/通用 seq2seqSeq2SeqTransformer训练返回 targets 的 logits推理走束搜索返回{outputs, scores}T5 风格 seq2seqT5Transformertf.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），仅供参考