基于deepset/gbert-base的二分类器训练尺寸不匹配问题求助

问题概述 #

基于deepset/gbert-base实现二分类器时，使用含德语文本列与0/1标签列的自定义数据集，设置num_labels=2并遵循Hugging Face官方教程操作，出现以下错误：

ValueError: Target size (torch.Size([8])) must be the same as input size (torch.Size([8, 2]))

已尝试的无效方案 #

• 确认PyTorch版本为2.0.0+cu118
• 确保标签为float类型且无空值
• 将标签形状从[0]/[1]改为[1,0]/[0,1]（one-hot格式）
• 通过Trainer API更换损失函数为nn.CrossEntropyLoss()
• 更换为nlptown/bert-base-multilingual-uncased-sentiment等其他模型

相关代码 #

from transformers import AutoTokenizer, DataCollatorWithPadding
tokenizer = AutoTokenizer.from_pretrained("deepset/gbert-base")

def tokenize_function(examples):
    return tokenizer(examples["text1"], truncation=True)

tokenized_datasets = final_dataset_dict.map(tokenize_function, batched=True)
data_collator= DataCollatorWithPadding(tokenizer)
tokenized_datasets = tokenized_datasets.remove_columns(["text1"])
tokenized_datasets["train"].column_names
tokenized_datasets.set_format("torch")

from torch.utils.data import DataLoader

train_dataloader = DataLoader(tokenized_datasets["train"], shuffle = True, batch_size = 8, collate_fn = data_collator)
eval_dataloader = DataLoader(tokenized_datasets["unsupervised"], batch_size = 8, collate_fn = data_collator)

for batch in train_dataloader:
  break
print({k: v.shape for k, v in batch.items()})

from transformers import AutoModelForSequenceClassification
checkpoint = "deepset/gbert-base"
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels =2)

outputs = model(**batch)
print(outputs.loss, outputs.logits.shape)

批量数据形状 #

{'labels': torch.Size([8]), 'input_ids': torch.Size([8, 69]), 'token_type_ids': torch.Size([8, 69]), 'attention_mask': torch.Size([8, 69])}

错误栈 #

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-36-b84c8f6552ab> in <cell line: 1>()
----> 1 outputs = model(**batch)
      2 #print(outputs.shape)
      3 print(outputs.loss, outputs.logits.shape)

4 frames
/usr/local/lib/python3.9/dist-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
   3161 
   3162     if not (target.size() == input.size()):
-> 3163         raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   3164 
   3165     return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)

ValueError: Target size (torch.Size([8])) must be the same as input size (torch.Size([8, 2]))

解决思路 #

1. 对齐任务类型、标签格式与损失函数 #

你设置了num_labels=2，模型会输出形状为[batch_size, 2]的logits（对应两个类别的预测得分），此时应该用类别索引格式的标签（形状[batch_size]）配合CrossEntropyLoss，而非BCE系列损失。错误的核心是模型自动选择了BCEWithLogitsLoss，导致要求标签和logits同形状。

2. 强制模型使用CrossEntropyLoss #

如果用Trainer API，需要自定义损失计算逻辑覆盖默认行为：

from torch.nn import CrossEntropyLoss
from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        # 计算CrossEntropyLoss，将logits和标签调整为匹配形状
        loss_fct = CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

之后用CustomTrainer替代原生Trainer即可。

3. 验证标签格式的正确性 #

• 不要将标签转为one-hot格式：CrossEntropyLoss要求标签是类别索引（0或1），而非[1,0]或[0,1]
• 确保标签维度为1：即每个样本对应一个整数标签，形状为[batch_size]，无需额外扩展维度

4. 检查模型任务类型的自动识别 #

Hugging Face的模型会根据数据集标签自动推断任务类型，如果你的标签曾被处理为浮点型，可能导致模型误判为二元分类（而非多分类）。可以显式在模型初始化时指定任务类型，或者确保标签为整数类型：

# 将标签转为整数类型
tokenized_datasets = tokenized_datasets.map(lambda x: {"labels": x["labels"].long()})

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