如何基于TF-Keras结合HuggingFace实现BERT额外预训练

云原生阿岚

2026-9-26

实现MLM预训练的后续步骤

你当前使用的TFDistilBertForMaskedLM已经内置了掩码语言建模的损失计算逻辑，只需要补充随机掩码处理和数据集构造两步即可启动训练，具体操作如下：

1. 导入依赖并配置数据整理器

直接使用HuggingFace内置的DataCollatorForLanguageModeling自动完成MLM所需的随机掩码、标签生成逻辑，无需手动实现掩码规则：

import tensorflow as tf
from transformers import DataCollatorForLanguageModeling, AdamWeightDecay

# 实例化MLM数据整理器，掩码概率设为标准0.15
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15,
    return_tensors="tf"
)

2. 优化器配置（可选但推荐）

原生adam优化器没有权重衰减，和BERT预训练的优化策略不符，建议替换为适配Transformers的AdamWeightDecay：

optimizer = AdamWeightDecay(learning_rate=2e-5, weight_decay_rate=0.01)
model.compile(optimizer=optimizer)

3. 构造训练数据集

将分词后的结果转为TensorFlow Dataset格式，应用数据整理器自动生成带掩码的输入和对应标签：

# 将分词结果转为TF数据集
dataset = tf.data.Dataset.from_tensor_slices(dict(data))
# 按批次处理，自动完成每批次的随机掩码
dataset = dataset.batch(2).map(lambda x: data_collator(x))

4. 启动训练

直接调用Keras的fit方法即可，HuggingFace的TF模型会自动识别输入中的labels字段作为训练目标计算损失：

model.fit(dataset, epochs=3)

补充说明

数据整理器生成的每批次数据包含三个字段：input_ids（带掩码的输入序列）、attention_mask（注意力掩码）、labels（训练标签，非掩码位置为-100，损失计算时自动忽略），完全匹配TFDistilBertForMaskedLM的输入要求，无需手动传参
若需要自定义掩码规则，可手动实现标准MLM逻辑：随机选择15%的token，80%替换为[MASK]、10%替换为随机token、10%保留原token，对应位置labels设为原token id、其余位置设为-100即可，逻辑和PyTorch实现完全一致，仅张量操作替换为TensorFlow API
额外预训练完成后，可通过model.save_pretrained("你的保存路径")存储权重，后续微调多标签分类任务时，用TFDistilBertForSequenceClassification.from_pretrained("你的保存路径")加载权重即可

内容的提问来源于stack exchange，提问作者Giuseppe Minardi