PyTorch将数据集移至CUDA后DataLoader迭代报错问题排查

错误原因 #

• DataLoader的核心采样逻辑（比如随机采样、批次划分）默认在CPU运行，这些组件需要读取数据索引或原始数据完成操作。当你把trainData.data和trainData.targets直接移到CUDA后，采样器尝试访问GPU张量时会触发自动转numpy操作——而GPU张量无法直接转为numpy（必须先移回CPU），因此抛出TypeError。
• 多数PyTorch内置Dataset的__getitem__方法默认基于CPU数据设计，会执行numpy转CPU张量的操作，当原始数据已是GPU张量时，会出现类型不兼容问题。

解决方法 #

方法1：自定义GPU兼容的Dataset类 #

直接在Dataset初始化时把数据全量移到GPU，重写__getitem__直接返回GPU张量，绕开CPU端的转换逻辑：

import torch
from torch.utils.data import Dataset, DataLoader

class CUDADataset(Dataset):
    def __init__(self, data, targets):
        # 初始化时直接把数据移到GPU
        self.data = data.to('cuda', non_blocking=True)
        self.targets = targets.to('cuda', non_blocking=True)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        # 直接返回GPU上的张量，无额外转换
        return self.data[idx], self.targets[idx]

# 从原始Dataset提取数据和标签
train_data = trainData.data
train_targets = trainData.targets

# 构建自定义GPU数据集
cuda_train_dataset = CUDADataset(train_data, train_targets)
# 注意：多进程无法访问主进程GPU张量，需设num_workers=0
train_loader = DataLoader(cuda_train_dataset, batch_size=32, shuffle=True, num_workers=0)

# 训练时直接使用，无需再手动移GPU
for data, target in train_loader:
    output = model(data)
    loss = criterion(output, target)
    # 后续优化步骤...

方法2：使用TensorDataset快速构建GPU数据集 #

PyTorch的TensorDataset原生支持GPU张量，无需自定义类：

from torch.utils.data import TensorDataset, DataLoader

# 把原始数据和标签转为CUDA张量
train_data_cuda = trainData.data.to('cuda', non_blocking=True)
train_targets_cuda = trainData.targets.to('cuda', non_blocking=True)

# 构建TensorDataset
train_dataset = TensorDataset(train_data_cuda, train_targets_cuda)

# 创建DataLoader，必须设num_workers=0
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=0)

# 训练迭代
for data, target in train_loader:
    # 数据已在GPU上，直接传入模型
    output = model(data)
    # ...训练逻辑

方法3：临时规避（不推荐） #

如果坚持用内置Dataset，需关闭多进程（num_workers=0），避免采样逻辑触发CPU-GPU转换：

# 把原始数据集的数据移到GPU
trainData.data = trainData.data.to('cuda')
trainData.targets = trainData.targets.to('cuda')

# 必须设置num_workers=0，否则多进程会报错
train_loader = DataLoader(trainData, batch_size=32, shuffle=True, num_workers=0)

# 训练迭代
for data, target in train_loader:
    # 数据已在GPU上
    output = model(data)
    # ...训练逻辑

注意事项 #

• 全量移GPU会占用大量显存，若数据集过大可能触发OOM（显存不足），需根据GPU显存容量判断是否适用。
• 只要用全量GPU数据集，必须设置num_workers=0——多进程工作者无法直接访问主进程的GPU张量，会导致数据复制错误。
• 内置Dataset的预处理逻辑（如归一化）可能依赖CPU操作，若原始数据是GPU张量，建议手动把预处理步骤移到GPU执行（比如用torch.nn.functional的GPU兼容函数）。

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