使用Transformers库训练文本分类模型时遇形状不兼容ValueError

错误根源 #

报错提示的Shapes (3, 1) and (512, 3)不兼容，本质是模型输出维度与标签维度/顺序完全不匹配：

• 模型输出(512, 3)：代表批次大小为512，对应3个分类类别的logits
• 标签输入(3, 1)：维度顺序颠倒，且多了一个不必要的维度，导致损失计算时无法对齐

具体修复步骤 #

1. 修正标签形状，去掉多余维度 #

确保标签是一维张量/数组（形状为(batch_size,)），而非二维的(batch_size,1)：

• 若用Hugging Face Dataset预处理数据，直接返回原始标签，不要额外扩展维度：

  def preprocess_data(examples):
      tokenized = tokenizer(
          examples["text"], 
          truncation=True, 
          padding="max_length",
          max_length=512
      )
      # 直接赋值原始label，不要用np.expand_dims或torch.unsqueeze
      tokenized["labels"] = examples["label"]
      return tokenized
  

• 如果标签已经是二维格式，在训练前用squeeze去除多余维度：

  # PyTorch示例
  labels = batch["labels"].squeeze(dim=1)
  # TensorFlow示例
  labels = tf.squeeze(labels, axis=1)
  

2. 匹配损失函数与标签格式 #

针对AutoModelForSequenceClassification的输出（(batch_size, num_classes)），选择对应损失函数：

• 标签为一维整数（如0、1、2）：用SparseCategoricalCrossentropy，必须开启from_logits=True：

  # PyTorch示例
  loss_fn = torch.nn.CrossEntropyLoss()  # 等价于SparseCrossEntropy，无需额外设置
  # TensorFlow示例
  loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
  

• 标签为one-hot编码的二维数组（如[1,0,0]）：用CategoricalCrossentropy：

  loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
  

3. 检查数据加载的批次维度 #

确认数据加载器返回的批次中，标签的第一维度是批次大小：

from torch.utils.data import DataLoader

train_loader = DataLoader(tokenized_train_dataset, batch_size=512, shuffle=True)
for batch in train_loader:
    # 正常标签形状应为(512,)
    print("Labels shape:", batch["labels"].shape)
    # 模型输出logits形状应为(512, 3)
    outputs = model(**batch)
    print("Logits shape:", outputs.logits.shape)

4. 排查Trainer API的配置问题 #

如果用Hugging Face Trainer训练，确保标签格式符合要求，无需手动指定损失（Trainer会自动匹配）：

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    per_device_train_batch_size=512,
    num_train_epochs=3,
    logging_steps=10
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_eval
)

额外排查点 #

• 确认数据集标签是0开始的连续整数（3类对应0、1、2），无非法值或超出范围的标签
• 检查tokenizer的max_length设置，确保输入序列长度统一，避免间接影响批次维度

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