使用Huggingface BERT分类时输入与标签批量大小不匹配报错

核心原因 #

BertForSequenceClassification默认适配单标签分类任务，其内置的cross_entropy损失函数要求标签为类别索引格式（形状为[batch_size]），而非独热编码格式（[batch_size, num_classes]）。你传入的(32,24)独热标签会被模型内部错误展平为768（32×24），导致与输入batch_size（32）不匹配，触发报错。

解决方案 #

根据你的任务类型选择对应方案：

方案1：单标签分类（每个样本仅属于一个类别） #

将独热编码标签转换为类别索引：

for i in train_dataloader:
    i = tuple(t.to(device) for t in i)
    # 将(32,24)的独热标签转为(32,)的类别索引
    labels = torch.argmax(i[2], dim=1)
    print(i[0].shape, i[1].shape, labels.shape)  # 此时labels.shape应为torch.Size([32])
    output = model(i[0], attention_mask=i[1], labels=labels)
    print(output.logits.shape)
    break

方案2：多标签分类（每个样本可属于多个类别） #

如果你的任务是多标签分类，需要在初始化模型时指定problem_type="multi_label_classification"，此时模型会自动使用BCEWithLogitsLoss，支持独热编码格式的标签：

from transformers import BertForSequenceClassification

model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=24,
    problem_type="multi_label_classification"  # 关键设置
)

之后直接传入(32,24)的标签即可，模型会正确计算损失，不会再出现batch_size不匹配的问题。

验证说明 #

• 单标签场景下，模型输出的logits形状为[32,24]，对应每个样本在24个类别上的得分，cross_entropy会自动匹配索引标签计算损失。
• 多标签场景下，logits形状同样为[32,24]，BCEWithLogitsLoss会逐类别计算二分类损失，最终平均得到总损失。

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