训练T5-small乌尔都语-英语翻译模型遇维度不匹配ValueError

我尝试训练乌尔都语到英语的翻译模型，但持续遇到问题，训练代码如下：

import pandas as pd
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from transformers import TrainingArguments, Trainer
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.nn.utils.clip_grad import clip_grad_norm_
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import numpy as np
import os
import random

# Set the seed for reproducibility
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)


# Load the datasets
# IMPORT REQUIRED DATASET
path = "/content/train_set.csv"
path_val = "/content/dev_set.csv"
path_test = "/content/test_ur.csv"

ds_train = pd.read_csv(path)
ds_val = pd.read_csv(path_val)
ds_test = pd.read_csv(path_test)
ds_train

ds_train=ds_train.dropna()
ds_val=ds_val.dropna()
ds_test=ds_test.dropna()


train_df = ds_train
val_df = ds_val
test_df = ds_test


# Create a custom dataset class
class UrduEnglishDataset(Dataset):
    def __init__(self, df, tokenizer, max_source_length, max_target_length):
        self.df = df
        self.tokenizer = tokenizer
        self.max_source_length = max_source_length
        self.max_target_length = max_target_length

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

    def __getitem__(self, idx):
        urdu_text = self.df.iloc[idx, 0]
        english_text = self.df.iloc[idx, 1]

        encoding = self.tokenizer.encode_plus(
            urdu_text,
            add_special_tokens=True,
            max_length=self.max_source_length,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )

        labels = self.tokenizer.encode_plus(
            english_text,
            add_special_tokens=True,
            max_length=self.max_target_length,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': labels['input_ids'].flatten(),
            'labels_attention_mask': labels['attention_mask'].flatten()
        }

# Set the model parameters
model_params = {
    "MODEL": "t5-small",
    "TRAIN_BATCH_SIZE": 2,# 4, # 8,
    "VALID_BATCH_SIZE":1, #1 2, #4,
    "TRAIN_EPOCHS": 5, ##3
    "VAL_EPOCHS": 1,
    "LEARNING_RATE": 1e-4,
    "MAX_SOURCE_TEXT_LENGTH": 64, #128, #256, # 512,
    "MAX_TARGET_TEXT_LENGTH": 64, #128, #/ 256, # 512,
    "SEED": 42,
    "GRAD_CLIP": 1.0,
    "PATIENCE": 5
}

# Set the device (TPU or GPU)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load the tokenizer
tokenizer = T5Tokenizer.from_pretrained(model_params["MODEL"])

# Create the datasets and data loaders
train_dataset = UrduEnglishDataset(train_df, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"], model_params["MAX_TARGET_TEXT_LENGTH"])
val_dataset = UrduEnglishDataset(val_df, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"], model_params["MAX_TARGET_TEXT_LENGTH"])
test_dataset = UrduEnglishDataset(test_df, tokenizer, model_params["MAX_SOURCE_TEXT_LENGTH"], model_params["MAX_TARGET_TEXT_LENGTH"])

train_loader = DataLoader(train_dataset, batch_size=model_params["TRAIN_BATCH_SIZE"], shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=model_params["VALID_BATCH_SIZE"], shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=model_params["VALID_BATCH_SIZE"], shuffle=False)

# Load the model
model = T5ForConditionalGeneration.from_pretrained(model_params["MODEL"])

# Set the training arguments
training_args = TrainingArguments(
    output_dir='results_t5small',
    num_train_epochs=model_params["TRAIN_EPOCHS"],
    per_device_train_batch_size=model_params["TRAIN_BATCH_SIZE"],
    per_device_eval_batch_size=model_params["VALID_BATCH_SIZE"],
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='logs',
    logging_steps=10,
    eval_strategy='steps',
    save_steps=500,
    eval_steps=500,
    load_best_model_at_end=True,
    save_total_limit=5,
    report_to='tensorboard',
    learning_rate=model_params["LEARNING_RATE"],
    fp16=True
)
def compute_metrics(pred, label_ids=None):
    # Convert the label_ids to a tensor
    label_ids = label_ids if label_ids is not None else torch.tensor(pred.label_ids)

    # Convert the predictions to a tensor
    predictions = torch.tensor(pred.predictions)

    # Ensure the predictions tensor has the expected sequence length
    if predictions.size(3) != model_params["MAX_TARGET_TEXT_LENGTH"]:
        # Reshape the predictions tensor to match the expected length
        predictions = predictions.reshape(predictions.size(0), predictions.size(1), model_params["MAX_TARGET_TEXT_LENGTH"])

    # Compute the accuracy
    accuracy = torch.sum(label_ids == predictions.argmax(-1)).item()

    # Return the accuracy
    return {"accuracy": accuracy}
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    compute_metrics=compute_metrics #*lambda pred: {"accuracy": torch.sum(pred.label_ids == pred.predictions.argmax(-1)).item()}
)

# Train the model
history = trainer.train()

完整错误信息 #

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-6-68426953da61> in <cell line: 58>()
     56 
     57 # Train the model
---> 58 history = trainer.train()

5 frames
/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1883                 hf_hub_utils.enable_progress_bars()
   1884         else:
-> 1885             return inner_training_loop(
   1886                 args=args,
   1887                 resume_from_checkpoint=resume_from_checkpoint,

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   2289                     self.control = self.callback_handler.on_step_end(args, self.state, self.control)
   2290 
-> 2291                     self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
   2292                 else:
   2293                     self.control = self.callback_handler.on_substep_end(args, self.state, self.control)

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in _maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
   2719         metrics = None
   2720         if self.control.should_evaluate:
-> 2721             metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
   2722             self._report_to_hp_search(trial, self.state.global_step, metrics)
   2723 

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in evaluate(self, eval_dataset, ignore_keys, metric_key_prefix)
   3570 
   3571         eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
-> 3572         output = eval_loop(
   3573             eval_dataloader,
   3574             description="Evaluation",

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in evaluation_loop(self, dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix)
   3852                 )
   3853             else:
-> 3854                 metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
   3855         elif metrics is None:
   3856             metrics = {}

<ipython-input-6-68426953da61> in compute_metrics(pred, label_ids)
     35 
     36     # Convert the predictions to a tensor
---> 37     predictions = torch.tensor(pred.predictions)
     38 
     39     # Ensure the predictions tensor has the expected sequence length

ValueError: expected sequence of length 32128 at dim 3 (got 512)

补充说明 #

我曾尝试将model_params["MAX_TARGET_TEXT_LENGTH"]设置为32128测试，但导致会话终止，无法解决问题。

问题分析与解决方案 #

核心问题 #

错误根源在compute_metrics函数的逻辑误解：

• T5模型的预测输出pred.predictions形状为[batch_size, sequence_length, vocab_size]，其中32128是T5-small的词汇表大小，而非目标序列长度。你错误地认为第三维度是目标序列长度（64），强行转张量时引发维度不匹配。
• 将MAX_TARGET_TEXT_LENGTH设为32128完全错误，这会让目标序列长度等于词汇表大小，直接耗尽显存导致会话终止。

修复步骤 #

1. 修正数据集类的标签处理
   T5模型默认用-100标记padding位置以忽略loss，需修改数据集类的__getitem__方法：

def __getitem__(self, idx):
    urdu_text = self.df.iloc[idx, 0]
    english_text = self.df.iloc[idx, 1]

    encoding = self.tokenizer.encode_plus(
        urdu_text,
        add_special_tokens=True,
        max_length=self.max_source_length,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt'
    )

    labels = self.tokenizer.encode_plus(
        english_text,
        add_special_tokens=True,
        max_length=self.max_target_length,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt'
    )
    # 将padding的token id转为-100，让模型忽略这些位置的loss
    labels_input_ids = labels['input_ids'].flatten()
    labels_input_ids[labels_input_ids == tokenizer.pad_token_id] = -100

    return {
        'input_ids': encoding['input_ids'].flatten(),
        'attention_mask': encoding['attention_mask'].flatten(),
        'labels': labels_input_ids,
        'labels_attention_mask': labels['attention_mask'].flatten()
    }

2. 重写compute_metrics函数
   翻译任务不适合用普通分类准确率，推荐用token级准确率或专业的BLEU指标：

• 方案一：token级准确率

def compute_metrics(pred):
    # 获取每个位置概率最大的token id
    pred_ids = pred.predictions.argmax(-1)
    label_ids = pred.label_ids
    # 过滤padding位置（-100）
    mask = label_ids != -100
    
    # 计算正确token数和总有效token数
    correct = ((pred_ids == label_ids) & mask).sum().item()
    total = mask.sum().item()
    accuracy = correct / total if total > 0 else 0.0
    
    return {"token_accuracy": accuracy}

• 方案二：BLEU翻译指标（更专业）
  先安装依赖：!pip install sacrebleu

from sacrebleu.metrics import BLEU

def compute_metrics(pred):
    bleu = BLEU()
    # 将预测id转为文本，跳过特殊token
    pred_texts = tokenizer.batch_decode(pred.predictions.argmax(-1), skip_special_tokens=True)
    # 将标签id转为文本，过滤-100的padding
    label_texts = []
    for label in pred.label_ids:
        valid_label = label[label != -100]
        label_texts.append(tokenizer.decode(valid_label, skip_special_tokens=True))
    
    # 计算BLEU分数
    score = bleu.corpus_score(pred_texts, [label_texts])
    return {"bleu_score": score.score}

3. 清理无效参数
   TrainingArguments中没有VAL_EPOCHS参数，直接删除该配置项。

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