训练BERT多分类模型时出现target与output形状不匹配错误

错误原因分析 #

报错提示target.shape=(None, 512)与output.shape=(None, 3)不匹配，结合你提供的批次张量信息（input_ids/attention_mask为(16,512)，labels为(16,)），核心问题是数据集的结构与模型输入输出的映射不匹配：你的数据集元素应为((input_ids, attention_mask), labels)，但当前结构可能是(input_ids, attention_mask, labels)，导致训练时模型误将attention_mask（形状512）当作标签传入，与模型输出的3分类结果形状冲突。

解决步骤 #

1. 调整数据集结构 #

将输入特征（input_ids和attention_mask）打包成一个元组，与标签配对，确保数据集的每个元素格式为(输入特征元组, 标签)：

# 定义映射函数调整数据集结构
def format_data(input_ids, attention_mask, labels):
    return (input_ids, attention_mask), labels

# 应用到训练和验证数据集
train_dataset = train_dataset.map(format_data)
val_dataset = val_dataset.map(format_data)

# 重新构建数据加载器
train_dataloader = train_dataset.shuffle(buffer_size=10000).batch(batch_size)
validation_dataloader = val_dataset.batch(batch_size)

2. 验证数据加载器结构 #

打印一个批次的内容，确认输入和标签的形状符合预期：

for batch in train_dataloader.take(1):
    inputs, labels = batch
    print(f"Input IDs shape: {inputs[0].shape}")
    print(f"Attention Mask shape: {inputs[1].shape}")
    print(f"Labels shape: {labels.shape}")

正常输出应为：

Input IDs shape: (16, 512)
Attention Mask shape: (16, 512)
Labels shape: (16,)

3. 确认损失函数与标签格式匹配 #

由于你的标签是整数形式（形状(16,)），需使用SparseCategoricalCrossentropy作为损失函数（无需将标签转为one-hot编码）：

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=2e-5),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    metrics=['accuracy']
)

若后续将标签转为one-hot编码（形状(16,3)），则替换为CategoricalCrossentropy。

4. 检查模型输入层定义 #

确保模型输入层正确接收两个输入特征：

from transformers import TFBertModel

# 定义输入层
input_ids = tf.keras.layers.Input(shape=(512,), dtype=tf.int32, name="input_ids")
attention_mask = tf.keras.layers.Input(shape=(512,), dtype=tf.int32, name="attention_mask")

# 加载BERT模型并取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>输出
bert_model = TFBertModel.from_pretrained('bert-base-uncased')
bert_output = bert_model(input_ids, attention_mask=attention_mask)[1]

# 定义分类输出层
output = tf.keras.layers.Dense(3, activation="softmax")(bert_output)

# 构建模型
model = tf.keras.Model(inputs=[input_ids, attention_mask], outputs=output)

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