求助：HuggingFace训练器中数据加载器GPU迁移失败问题

问题 #

运行代码时触发错误：Expected all tensors to be on the same device, found two: cpu and cuda:0。已确认模型部署在cuda:0，但数据加载器输出的张量仍在CPU上，使用HuggingFace Transformers和Datasets库。

报错发生在继承自Seq2SeqTrainer的类的evaluate()方法中，相关代码段：

eval_dataset = self.eval_dataset if eval_dataset is None else eval_dataset
eval_dataloader = self.get_eval_dataloader(eval_dataset)
eval_examples = self.eval_examples if eval_examples is None else eval_examples
compute_metrics = self.compute_metrics
self.compute_metrics = None
eval_loop = (self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop)
try:
    # 报错位置
    output = eval_loop(
        eval_dataloader,
        description="Evaluation",
        prediction_loss_only=True if compute_metrics is None else None,
        ignore_keys=ignore_keys,
    )

尝试手动迁移张量时触发Too many values to unpack (expected 2)错误，代码：

for i, (inputs, labels) in eval_dataloader:
    inputs, labels = inputs.to(device), labels.to(device)

需解决：如何将DataLoader的批次移至GPU？能否通过修改Trainer的evaluation_loop实现？

解决方法 #

1. 预处理阶段绑定设备 #

在数据集加载完成后，通过set_format直接指定设备，让DataLoader生成的张量默认在GPU上：

dataset.set_format("torch", device="cuda:0")

或者在自定义预处理函数中，将处理后的张量移至GPU：

def preprocess_function(examples):
    processed = tokenizer(examples["text"], truncation=True, padding="max_length")
    for key in processed:
        processed[key] = torch.tensor(processed[key]).to("cuda:0")
    return processed

2. 重写Trainer的evaluation_loop #

继承Seq2SeqTrainer后，在evaluation_loop中自动迁移批次到模型所在设备：

from transformers import Seq2SeqTrainer
from torch.utils.data import DataLoader

class CustomSeq2SeqTrainer(Seq2SeqTrainer):
    def evaluation_loop(
        self,
        dataloader,
        description,
        prediction_loss_only=None,
        ignore_keys=None,
        metric_key_prefix="eval",
    ):
        device = self.model.device
        # 转换所有批次到目标设备
        processed_batches = []
        for batch in dataloader:
            processed_batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
            processed_batches.append(processed_batch)
        # 重构DataLoader
        new_dataloader = DataLoader(processed_batches, batch_size=dataloader.batch_size, shuffle=False)
        # 调用父类逻辑
        return super().evaluation_loop(
            new_dataloader,
            description,
            prediction_loss_only,
            ignore_keys,
            metric_key_prefix,
        )

使用该自定义Trainer替代原类即可自动处理设备迁移。

3. 修正手动遍历DataLoader的代码 #

HuggingFace DataLoader返回的批次是字典格式，而非(inputs, labels)元组，正确的迁移方式：

device = torch.device("cuda:0")
for batch in eval_dataloader:
    batch = {k: v.to(device) for k, v in batch.items()}
    # 后续可通过batch["input_ids"]、batch["labels"]访问对应张量

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