使用HuggingFace预训练Aleph-BERT做电影分类时Batch Size不匹配报错

错误核心原因 #

你用错了模型：BertForTokenClassification是做token级分类（比如命名实体识别）的，会输出每个输入token的预测结果（形状为[batch_size, sequence_length, num_labels]），但你的任务是电影流派文本分类（句子级单标签分类），标签形状是[batch_size]，两者维度完全不匹配，导致交叉熵计算时出现输入与目标batch_size不匹配的错误。

具体修复步骤 #

1. 替换为正确的文本分类模型 #

把BertForTokenClassification换成专门处理句子级分类的BertForSequenceClassification：

from transformers import BertForSequenceClassification

# 替换原模型定义行
model = BertForSequenceClassification.from_pretrained("onlplab/alephbert-base", num_labels=2)

2. 修正标签处理逻辑 #

原代码中标签处理存在两处问题：

• inplace=True会直接修改原DataFrame，且rename返回None，导致后续y值异常；
• 分类标签应使用整数类型而非浮点型。
  修改为：

y = movies_df[genre].astype(int).rename('labels')

3. 修复Trainer类的方法定义 #

你的MyTrainer类中compute_metrics方法缺少self参数，会导致调用报错，修正后：

class MyTrainer(Trainer):  
    def compute_metrics(self, pred):
        labels = pred.label_ids
        preds = pred.predictions.argmax(-1)
        precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
        acc = accuracy_score(labels, preds)
        return {
            'accuracy': acc,
            'f1': f1,
            'precision': precision,
            'recall': recall
        }

4. 确认数据集类的正确性 #

确保TextData类正确返回样本数据，示例实现（如果你的类未正确编写）：

from torch.utils.data import Dataset

class TextData(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

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