使用Hugging Face Roberta做多分类时遇Target Size不匹配错误求助

问题重现 #

使用Hugging Face Roberta进行15分类任务时触发以下错误：

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

用户代码 #

from datasets import load_dataset
from transformers import RobertaTokenizerFast, Trainer, DataCollatorWithPadding


dataset = load_dataset('csv', data_files=data_path,split = 'train')
train_testvalid = dataset.train_test_split(test_size=0.2)
test_valid = train_testvalid['test'].train_test_split(test_size=0.2)

checkpoint = "roberta-base"
tokenizer = RobertaTokenizerFast.from_pretrained(checkpoint)


def tokenization(example):
    return tokenizer(example['text'], truncation=True,max_length = 256, padding = True)


train_data = test_valid['train']
test_data= test_valid['test']

train_data = train_data.map(tokenization, batched = True)
test_data = test_data.map(tokenization, batched = True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)


train_data = train_data.remove_columns(["Unnamed: 0", "text"])
test_data = test_data.remove_columns(["Unnamed: 0", "text"])
train_data.set_format("torch")
test_data.set_format("torch")
train_data.column_names


from transformers import TrainingArguments

training_args = TrainingArguments("test-trainer")


from transformers import RobertaForSequenceClassification

model = RobertaForSequenceClassification.from_pretrained(checkpoint, num_labels=15)


from sklearn.metrics import accuracy_score, precision_recall_fscore_support
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='micro')
    acc = accuracy_score(labels, preds)
    return {
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall
    }


from transformers import Trainer

trainer = Trainer(
    model,
    training_args,
    compute_metrics=compute_metrics,
    train_dataset=train_data,
    eval_dataset=test_data,
    data_collator=data_collator,
    tokenizer=tokenizer,
)

trainer.train()

错误堆栈 #

ValueError                                Traceback (most recent call last)

<ipython-input-9-3435b262f1ae> in <module>
----> 1 trainer.train()

8 frames

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
   3146 
   3147     if not (target.size() == input.size()):
-> 3148         raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   3149 
   3150     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, 15]))

解决方案 #

核心原因 #

模型误将任务识别为多标签分类（使用二元交叉熵损失），但你的数据是单标签多分类（标签为一维类别索引），导致损失函数维度不匹配。

具体修复步骤 #

1. 确保标签列名为labels
   Hugging Face Trainer默认会读取名为labels的列作为目标。如果你的数据集标签列是其他名称（比如label），必须重命名：

   # 替换"your_label_column"为实际标签列名
   train_data = train_data.rename_column("your_label_column", "labels")
   test_data = test_data.rename_column("your_label_column", "labels")
   

2. 手动指定损失计算逻辑（可选但更稳妥）
   在初始化Trainer时，显式使用模型内置的损失计算，避免自动选择错误的损失函数：

   trainer = Trainer(
       model,
       training_args,
       compute_metrics=compute_metrics,
       train_dataset=train_data,
       eval_dataset=test_data,
       data_collator=data_collator,
       tokenizer=tokenizer,
       # 显式调用模型的损失计算，适配多分类任务
       compute_loss=lambda model, inputs: model(**inputs).loss
   )
   

3. 验证标签格式
   确认你的标签是一维张量（如torch.Size([8])），这是单标签多分类的正确格式，模型输出的[8,15]会被交叉熵损失自动处理维度匹配。

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