1. 从一次数据加载卡顿说起为什么需要搞懂这三个类刚接触深度学习那会儿我写过一个图像分类的小项目。数据集不大几千张图我图省事把所有图片一次性读进内存转成张量然后手动写了个循环每个batch切一刀喂给模型。训练倒是能跑但每次改个batch size或者想加个打乱顺序的功能就得回去改循环里的逻辑改来改去代码越来越乱。更难受的是后来换了个稍大一点的数据集内存直接爆了程序还没开始训练就崩了。那次之后我才认真去看了PyTorch的数据加载体系发现TensorDataset、Dataset、DataLoader这三个东西其实是一套组合拳各司其职。Dataset是抽象基类定义了数据从哪来、怎么取一条的规范TensorDataset是官方给的一个现成实现专门用来包装已经转成张量的数据DataLoader则负责调度把Dataset里的数据按batch取出来、打乱顺序、并行加载。三者配合起来代码干净内存可控扩展也方便。这篇内容就是把我踩过的坑和后来总结的用法梳理一遍。不管你是刚学PyTorch的新手还是写过几个项目但一直用土办法加载数据的人看完应该能把这套体系用顺。我会从最基础的概念讲起然后给出可直接复现的代码再聊几个实际项目中容易忽略的细节最后说说什么时候该自己写Dataset、什么时候用TensorDataset就够了。2. 三个角色的分工Dataset是规范TensorDataset是快捷实现DataLoader是调度器2.1 Dataset到底定义了什么PyTorch里的Dataset是一个抽象类它本身不能直接实例化使用作用是给所有数据集定义一个统一的接口。你继承它之后必须实现两个方法__len__和__getitem__。__len__返回数据集里有多少条样本DataLoader靠它知道一共要取多少数据、一个epoch有多少个batch。__getitem__接收一个索引返回一条样本这条样本可以是张量、可以是元组、也可以是字典具体格式由你自己定。为什么要有这么一层抽象因为实际项目里数据来源五花八门有的是硬盘上的图片文件夹有的是CSV文件有的是数据库查询结果还有的是实时生成的合成数据。如果每个地方都自己写加载逻辑代码风格会非常混乱。有了Dataset这个统一接口上层的数据调度逻辑就不用关心底层数据长什么样只管按索引取就行。我个人的理解是Dataset解决的是数据怎么读的问题它把读取逻辑封装起来让训练循环变得干净。你想想如果训练代码里混着文件路径拼接、图像解码、归一化这些操作那代码根本没法维护。2.2 TensorDataset帮你省掉重复劳动TensorDataset是Dataset的一个子类官方已经帮你实现好了__len__和__getitem__。它的用法很简单把几个张量传进去它会按第一个维度对齐每次取索引时返回这些张量在对应位置的切片。举个例子你有特征张量X形状是(1000, 20)标签张量y形状是(1000,)。用TensorDataset(X, y)包一下取第i条时就会返回(X[i], y[i])。就这么直接。它适合什么场景数据已经全部在内存里、已经转成张量、不需要复杂的预处理。比如你从numpy数组转过来的数据、从预训练特征提取器输出的特征、或者做小规模实验时手动构造的数据。这种情况下用TensorDataset最省事不用自己写类。但它也有明显的局限。第一它要求所有数据都能放进内存数据量大了就不行。第二它不支持自定义的预处理逻辑比如你想在取数据时做随机裁剪、加噪声TensorDataset做不到。第三它返回的是固定顺序的元组如果你想要字典格式的输出得自己再包一层。2.3 DataLoader才是真正干活的那个DataLoader接收一个Dataset对象然后负责按batch size把数据分组、决定是否打乱顺序、决定用几个进程加载、决定怎么把多条样本拼成一个batch。它内部的工作流程大致是这样的先根据__len__知道总共有多少条然后按batch_size算出有多少个batch。如果shuffleTrue就生成一个随机排列的索引列表。然后按索引去Dataset里取数据取够一个batch就通过collate_fn拼起来返回。collate_fn是个容易被忽略但很重要的参数。默认情况下它会把一个batch里的样本按位置堆叠成张量。比如每条样本是(特征, 标签)一个batch有32条默认就会返回两个张量形状分别是(32, 特征维度)和(32,)。但如果你的样本长度不一致比如文本序列默认的collate就会报错这时候就需要自定义collate_fn来做padding。还有一个参数num_workers控制用几个子进程加载数据。设成0表示在主进程里加载设成大于0会开子进程。这个参数对训练速度影响很大但也不是越大越好后面会细说。3. 从零写一个可运行的例子TensorDataset配DataLoader的完整流程3.1 构造数据与包装先来看一段可以直接跑的代码。假设我们做一个回归任务输入是20维特征输出是一个连续值。import torch from torch.utils.data import TensorDataset, DataLoader # 构造模拟数据 torch.manual_seed(42) N 1000 X torch.randn(N, 20) true_w torch.randn(20, 1) y X true_w 0.1 * torch.randn(N, 1) # 包装成TensorDataset dataset TensorDataset(X, y) print(len(dataset)) # 1000 print(dataset[0]) # (X[0], y[0]) 两个张量这里dataset[0]返回的是一个元组第一个元素是20维特征第二个元素是1维标签。注意标签的形状是(1,)而不是标量因为原始y是(N, 1)切片后保持了这个维度。3.2 DataLoader的配置与迭代dataloader DataLoader( dataset, batch_size32, shuffleTrue, num_workers0, drop_lastFalse ) for batch_X, batch_y in dataloader: print(batch_X.shape, batch_y.shape) break # 输出: torch.Size([32, 20]) torch.Size([32, 1])这段代码里几个参数值得说一下。batch_size32表示每个batch取32条。shuffleTrue表示每个epoch开始前重新打乱索引顺序这对训练很重要能防止模型学到样本顺序带来的虚假规律。drop_lastFalse表示如果最后一批不够32条也保留如果设成True就会丢掉最后那批不足的。num_workers0表示不用子进程数据在主进程里加载。对于这种全在内存里的张量数据设成0反而更快因为开子进程有额外的进程创建和数据传输开销。3.3 训练循环里怎么用import torch.nn as nn model nn.Sequential( nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 1) ) optimizer torch.optim.Adam(model.parameters(), lr1e-3) loss_fn nn.MSELoss() for epoch in range(5): model.train() total_loss 0.0 for batch_X, batch_y in dataloader: optimizer.zero_grad() pred model(batch_X) loss loss_fn(pred, batch_y) loss.backward() optimizer.step() total_loss loss.item() * batch_X.size(0) print(fepoch {epoch}, avg loss {total_loss / len(dataset):.4f})这个训练循环就是标准写法。每次从dataloader里取出一个batch前向、算损失、反向、更新。注意loss.item() * batch_X.size(0)是为了按样本数加权平均因为最后一个batch可能不满。跑下来你会发现用DataLoader之后训练代码里完全没有数据加载的细节全是模型相关的逻辑。这就是这套体系的价值关注点分离。4. 当TensorDataset不够用自定义Dataset的典型场景与写法4.1 什么时候必须自己写DatasetTensorDataset虽然方便但下面这些情况它搞不定数据在硬盘上不能一次性读进内存比如几万张高分辨率图片。每条样本需要不同的预处理比如随机裁剪、随机翻转、加噪声。数据来源是数据库或网络接口需要按需查询。样本长度不一致需要动态padding。需要返回字典格式而不是元组方便后续处理。这些场景下就得继承Dataset自己实现。核心还是那两个方法但__getitem__里可以写任意逻辑。4.2 一个图像数据集的完整实现假设我们有一个图片文件夹每张图对应一个类别标签标签存在一个CSV文件里。下面是一个简化但可运行的实现import os import csv import torch from torch.utils.data import Dataset from PIL import Image import torchvision.transforms as T class ImageFolderDataset(Dataset): def __init__(self, root_dir, label_file, transformNone): self.root_dir root_dir self.transform transform self.samples [] with open(label_file, r) as f: reader csv.reader(f) next(reader) # 跳过表头 for row in reader: img_name, label row[0], int(row[1]) self.samples.append((img_name, label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_name, label self.samples[idx] img_path os.path.join(self.root_dir, img_name) image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, label这个类里__init__只做轻量的事读标签文件、存路径列表。真正的图片读取和解码放在__getitem__里这样只有被取到的样本才会加载内存占用可控。transform参数用来传入预处理管道比如train_transform T.Compose([ T.RandomResizedCrop(224), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])注意ToTensor()必须放在归一化之前因为Normalize只接受张量。这个顺序搞反了会报错我见过不止一个新手在这里卡住。4.3 自定义Dataset配合DataLoader的注意事项自定义Dataset和DataLoader配合时有几个点容易出问题。第一__getitem__里不要修改类的成员变量。因为如果num_workers 0每个子进程会复制一份Dataset对象你在子进程里改的东西主进程看不到而且多个worker之间也不共享。所有需要随机的操作比如随机裁剪应该用transform里的随机函数它们内部有自己的随机状态。第二如果__getitem__里有可能抛异常比如图片损坏、文件缺失最好在内部捕获并返回一个替代样本否则整个训练会崩。我一般会加个try-except遇到坏样本就返回一个全零张量和忽略标签或者直接随机换一个索引重取。第三__len__返回的长度必须和实际可取的索引范围一致。如果返回的长度比实际样本多DataLoader取到后面会报IndexError。5. DataLoader那些影响训练效率和结果的参数5.1 num_workers设多少合适num_workers是最常被问到的参数。它的作用是开多少个子进程来并行加载数据。设成0表示不用子进程在主进程里串行加载。什么时候该设大于0当数据加载成为瓶颈的时候。比如你从硬盘读图片、做复杂的图像增强这些操作耗时较长如果串行做GPU就会闲着等数据。这时候开几个worker并行加载能让GPU一直有活干。但也不是越大越好。每个worker都是一个独立的进程会占用内存和CPU。worker太多进程间切换的开销反而拖慢速度。经验值是设成CPU核心数的一半左右比如8核机器设4。但最终还是要实测用time模块测一下一个epoch的耗时对比不同设置。还有一个坑在Windows上num_workers 0时DataLoader的构造必须放在if __name__ __main__:保护块里否则会无限递归创建进程。这个在Linux上不是问题但跨平台代码要注意。5.2 shuffle、drop_last与sampler的关系shuffleTrue本质上是在每个epoch开始时用RandomSampler替换默认的SequentialSampler。如果你传了自定义的sampler参数那shuffle就必须是False否则会报冲突。drop_last决定是否丢掉最后一个不满的batch。什么时候该设True当你的模型里有BatchNorm层而最后一个batch只有1条样本时BatchNorm会报错因为它在训练时需要至少2条样本来计算方差。这种情况下设drop_lastTrue能避免报错。但丢掉最后一批也意味着损失了一部分数据。如果数据量本来就少丢掉可能影响训练效果。折中办法是把batch size调小一点让最后一批不至于只有1条。5.3 collate_fn的默认行为与自定义时机默认的collate_fn会把一个batch里的样本按结构堆叠。如果每条样本是(张量A, 张量B)它会返回(堆叠后的A, 堆叠后的B)。堆叠要求所有样本在对应位置上的形状一致。什么时候需要自定义最常见的是变长序列。比如每条样本是一个不同长度的文本序列默认collate会因为形状不一致而报错。这时候需要写一个函数把序列padding到batch内最大长度。def pad_collate(batch): sequences, labels zip(*batch) max_len max(seq.size(0) for seq in sequences) padded torch.zeros(len(sequences), max_len, dtypetorch.long) for i, seq in enumerate(sequences): padded[i, :seq.size(0)] seq return padded, torch.tensor(labels)这个函数接收一个列表每个元素是__getitem__返回的样本。先解包算出最大长度然后创建全零张量把每条序列填进去。最后返回padding后的张量和标签张量。自定义collate_fn时要注意它是在主进程里执行的不在worker里。所以如果padding逻辑很重会成为瓶颈。一般padding操作本身很快问题不大。6. 几个我实际踩过的坑和对应的排查思路6.1 内存暴涨worker进程的数据复制有一次我做一个较大的数据集num_workers设了8结果内存直接爆了。排查后发现每个worker进程都会复制一份Dataset对象。如果Dataset的__init__里加载了大量数据到内存那8个worker就是8份内存直接乘以8。解决办法是把重数据的加载推迟到__getitem__里__init__只存路径和元信息。这样每个worker只在需要时才读具体样本不会一开始就复制全部数据。另一个相关的问题是如果__getitem__里缓存了数据比如用一个字典存已读过的图片那每个worker都有自己的缓存主进程看不到也起不到共享作用。这种缓存只在单个worker内部有效跨worker不共享。6.2 训练结果不可复现随机种子的设置用DataLoader时如果shuffleTrue且num_workers 0即使你设了torch.manual_seed每次运行的结果也可能不一样。原因是每个worker有自己的随机状态主进程的种子不会自动传给它们。要完全复现需要给worker设置种子。PyTorch提供了一个worker_init_fn参数def worker_init_fn(worker_id): np.random.seed(42 worker_id) torch.manual_seed(42 worker_id) dataloader DataLoader(dataset, batch_size32, shuffleTrue, num_workers4, worker_init_fnworker_init_fn)这样每个worker的随机种子是确定的配合主进程的种子就能复现。不过要注意即使这样如果用了CUDA的某些非确定性操作还是可能不完全一致。完全复现是个系统工程这里只是解决数据加载部分的随机性。6.3 最后一个batch引发的形状错误前面提过BatchNorm在单样本batch时会报错。还有一种情况是模型里用了view或reshape而最后一个batch的样本数不是预期值导致形状对不上。比如模型里写了x.view(batch_size, -1)但最后一个batch实际只有16条而batch_size设的是32就会报错。正确写法是用x.view(x.size(0), -1)让形状根据实际数据来。这个坑的本质是不要假设每个batch的大小都等于设定的batch_size。要么用drop_lastTrue保证要么在代码里用动态形状。6.4 数据加载成为瓶颈的识别方法怎么判断数据加载是不是瓶颈一个简单办法是在训练循环里计时。如果GPU利用率低而CPU占用高大概率是数据加载拖后腿。更精确的做法是分别测数据加载时间和模型前向反向时间import time # 测数据加载 t0 time.time() for _ in dataloader: pass t_data time.time() - t0 # 测模型计算 batch next(iter(dataloader)) t0 time.time() for _ in range(100): pred model(batch[0]) loss loss_fn(pred, batch[1]) loss.backward() t_model time.time() - t0如果t_data和t_model接近甚至更大说明数据加载是瓶颈可以考虑增加num_workers、简化预处理、或者把数据预处理提前做好存成缓存文件。7. 从TensorDataset到自定义Dataset的平滑过渡建议如果你现在用的是TensorDataset但发现数据量变大或者需要加预处理想迁移到自定义Dataset我建议按这个顺序来。第一步先把TensorDataset替换成一个最简单的自定义Dataset__getitem__里直接返回原来TensorDataset返回的东西。跑一遍确认训练结果和之前一致说明接口对接没问题。第二步把数据加载逻辑从__init__里挪到__getitem__里。比如原来是__init__里读所有图片现在改成__init__只存路径__getitem__里读单张。这一步做完内存占用应该会明显下降。第三步加入预处理逻辑。把transform管道加到__getitem__里注意训练和验证要用不同的transform。验证集一般不做随机增强只做resize和归一化。第四步调整DataLoader参数。数据从内存变成硬盘读取后加载耗时增加这时候把num_workers调大观察训练速度变化。这个渐进式迁移的好处是每一步都可验证出了问题容易定位是哪一步引入的。直接重写整个数据管道然后调试往往会被各种小问题缠住。另外说一个我个人的习惯不管用TensorDataset还是自定义Dataset我都会先写一个小的测试脚本单独把Dataset和DataLoader拿出来跑一遍打印几个batch的形状和内容确认没问题了再接到训练循环里。这样能把数据问题和模型问题分开排查省很多时间。最后提一下PyTorch还有IterableDataset这个类适合流式数据或者数据量大到无法用索引访问的场景。但那是另一个话题了对于大多数项目Dataset加DataLoader这套组合已经够用。把这两个用熟再根据实际需求决定要不要上更复杂的方案。