Zero基础学PyTorch前面几篇已经过完了张量、自动求导、模型搭建这些基础。到了第五篇咱们进入一个非常实际、也最容易让人卡住的话题DataLoader。标题里写的是“Dataloder”估计是少写了个a它说的就是PyTorch的数据加载器。很多人一上来喜欢把整个数据集塞进内存或者用一个循环手动取数据结果要么内存爆炸要么训练速度惨不忍睹。这篇我打算把DataLoader从底层机制讲到实际踩坑再把常被忽视的nn.Sequential、nn.ModuleList、nn.ModuleDict这些容器模块串起来讲清楚。内容尽量零基础代码可以直接跑。1. 从“一次全塞进内存”的教训聊起DataLoader到底解决什么问题1.1 为什么不能一次把所有数据都转成Tensor先想一个场景你手上有一套图像数据集几十万张图片每张几百KB。很多人第一次写训练脚本下意识就是遍历文件夹把每张图都读进来resize、归一化然后堆成一个巨大的Tensor。训练的时候整个模型吃进去看起来特别爽。但问题马上就会出现。直观说一张224x224的RGB图像转成浮点Tensor之后就是3×224×224×4字节大概600多KB。你以为几十万张图只是“几十万个小文件”换算完之后发现光这一个Tensor的内存占用就超过100GB。还没开始训练机器先挂掉了。这是内存问题。另一个问题更隐蔽叫做“随机梯度下降的随机性”。如果你把整个数据集一次性加载进来训练时按顺序把每个样本喂给模型那每个batch的样本分布是高度相关的——前面一批全是某一类后面一批又是另一类。梯度震荡会非常剧烈甚至收敛不到好的结果。DataLoader那个shuffleTrue可不是可有可无的装饰它直接影响训练效果。第三个问题是“计算模式”。数据读入、解码、缩放、归一化这些操作在GPU上做很浪费在CPU上做又会让训练循环变得特别啰嗦。你当然可以在训练代码里写一堆预处理逻辑但那个代码会越来越长到最后你自己都看不懂。DataLoader做的事就是把这套“读数据-预处理-分批-打乱”的流程封装成一条流水线你只需要在训练循环里说一句for batch in loader它就把一切准备好了。1.2 DataLoader把训练数据流程化成一条流水线拿生产车间打个比方。数据集是仓库里的原材料DataLoader是传送带模型是组装工位。如果没有传送带组装工就得自己一趟趟跑去仓库搬货搬回来还要自己拆包装、质检、分堆。DataLoader这根传送带的意义在于它帮你把原料从仓库里拉出来按固定节奏送到工位前每次送多少个batch_size、送哪批货shuffle、要不要拆开检验transform、用几个人搬货num_workers都一次性配置好。在代码层面DataLoader是围绕torch.utils.data.DataLoader这个类来工作的。它接收一个Dataset对象作为数据源然后负责三件事样本组织按batch大小切分组决定要不要打乱顺序多进程预取开几个子进程提前把数据读好训练的时候不用干等I/O批次拼装把若干个样本的Tensor拼接成一个批量Tensor这一步就是collate_fn在做。所以DataLoader本身并不存储数据它只是数据工厂里的那个调度系统。真正的数据源头在Dataset里。这也是为什么大家总是把Dataset和DataLoader拿出来一起讲——两者是前后端的关系。认识到这个分工之后写起代码就清楚多了你自己写一个继承torch.utils.data.Dataset的类定义好“如何取出一条样本”剩下的分批、打乱、并行读取全交给DataLoader。2. 先定义一个自己的 Dataset理解数据接口的约定2.1 自定义Dataset的骨架Dataset在PyTorch里是一个抽象接口要求子类实现两个魔法方法__len__和__getitem__。前者返回样本总数后者根据索引返回一条样本。就这么简单。很多人去读官方文档看得头大其实内部就这两件事。举个例子假设我们有一个回归任务数据是一段正弦波信号每个样本是一个128长度的向量标签是它的频率。你可以这样写import torch from torch.utils.data import Dataset class SineWaveDataset(Dataset): def __init__(self, num_samples10000, seq_len128): self.num_samples num_samples self.seq_len seq_len # 预先生成样本和标签 self.samples torch.randn(num_samples, seq_len) t torch.arange(seq_len).float() self.labels torch.sin(t).unsqueeze(0).repeat(num_samples, 1) def __len__(self): # 返回样本数量 return self.num_samples def __getitem__(self, idx): # 按索引取出第idx条样本 return self.samples[idx], self.labels[idx]这段代码里有几个关键点。第一__init__里可以把数据常驻内存也可以只存文件路径列表取决于数据规模。第二__len__必须返回一个整数PyTorch会在shuffle、划分数据集时用到它。第三__getitem__通常返回一个tuple前一个元素是输入后一个是标签。DataLoader会把同一个batch里返回的第一项拼成一个Tensor第二项拼成另一个Tensor。2.2 getitem里到底该做什么、不该做什么__getitem__是数据管线的核心很多人在这里犯错误。第一种错误是把整个数据集的预处理放到__getitem__里做比如每次按索引去读文件、解码、resize。这本身没错但要注意它会被频繁调用。训练一个epoch要调用num_samples次如果你的预处理太重训练时间会成倍拉长。第二种错误是在__getitem__里做随机打乱。有人写shuffleTrue之后还在这个函数里对样本做随机旋转、裁剪这没问题属于数据增强。但如果你把整个数据集的索引顺序也放在里面打乱那就没意义了——打乱索引是DataLoader在外部控制的你内部改它只会造成混乱。合理的分工是__init__里做一次性的数据读取或路径扫描__getitem__里做单条样本的读取和轻量预处理批量层面的操作交给DataLoader和collate_fn。如果某个预处理特别耗时最好先离线跑一遍把结果存成预处理后的文件训练时直接读这叫“离线增强和在线增强的权衡”。2.3 用transform把预处理从训练循环里剥离torchvision.transforms这套工具很多人不太理解为什么要单独拆一层。其实它的好处是把图像的预处理动作串联起来变成一个可组合、可复用的对象。比如from torchvision import transforms transform transforms.Compose([ transforms.Resize((224, 224)), # 缩放 transforms.RandomHorizontalFlip(), # 随机水平翻转 transforms.ToTensor(), # 转成Tensor像素值归一化到0-1 transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), # 标准化 ])然后在Dataset里加一个self.transform参数在__getitem__里调用它x self.transform(x)。这样一来训练集的预处理和验证集的预处理可以复用同一个Compose只是在训练集上把RandomHorizontalFlip去掉就很干净。这里有一个零基础同学容易忽略的细节ToTensor()会把像素值从0-255变成0-1Normalize再把它变成均值为0、方差为1的分布。这两步顺序不能乱因为归一化的均值和标准差是为0-1范围的像素设计的。你要是先Normalize再ToTensor数值范围完全不对。模型在推理阶段拿到的数据分布必须和训练阶段一致这是很多模型效果莫名变差的隐藏原因。3. DataLoader参数逐个拆解每一个都对应一个问题定义好Dataset之后把它交给DataLoader。这一节把常用参数按“它们解决什么问题”来拆比死记参数表管用得多。from torch.utils.data import DataLoader loader DataLoader( datasetSineWaveDataset(num_samples10000), batch_size32, # 每个批次多少条样本 shuffleTrue, # 每个epoch进来是否打乱顺序 num_workers4, # 用几个子进程准备数据 pin_memoryTrue, # 分配给GPU时锁页内存 drop_lastTrue, # 末尾不足一批时是否丢弃 collate_fnNone, # 自定义批次拼接函数 )3.1 batch_size与drop_lastbatch_size是每个批次包含的样本数这个参数直接影响显存占用和梯度更新频率。batch太大显存可能爆掉batch太小梯度噪声大训练不稳定。到底设多少合适跟模型大小、显卡显存都有关系。零基础可以先从32或64开始遇到OOM就往下减比如16、8训练loss下降太飘就往上加一点。drop_last处理的是最后不足一个batch的情况。假设你有10000条样本batch_size是128那最后一个batch只有16条。不丢弃的话这个batch的统计分布和其他batch差别很大可能会让BN层BatchNorm的统计量被带偏。所以训练阶段我通常设drop_lastTrue验证阶段设drop_lastFalse这样验证指标能覆盖到所有样本。3.2 shuffle打乱顺序背后的梯度含义shuffle的作用前面已经提到过这里再往深说一层。深度学习实际用的是小批量随机梯度下降每一轮迭代用一个batch来估计整个数据集的梯度方向。如果数据顺序不随机某个batch可能全是简单样本另一个batch全是难样本梯度方向就会来回摆。让每个batch尽可能代表整个数据集的分布是shuffle的核心目的。尤其在分类任务里如果你的数据文件恰好按类别排序不shuffle的话先训练到的是全是第0类后面全是第1类模型会频繁遗忘前面的东西。很多刚开始练手的同学说“为什么我loss在下降但准确率上不去”检查一下大概率是忘了开shuffle。3.3 num_workers与pin_memorynum_workers是让DataLoader用几个子进程提前准备数据。设成0表示只在主进程里同步加载设成4表示开4个子进程并行读数据。这个参数最大的价值是“预取”——训练在GPU上算当前batch的时候子进程已经在读下一个batch了计算和I/O重叠起来训练速度会明显提升。但这不是越大越好。每个worker都会复制一部分内存开太多可能内存不足还会增加进程通信开销。一个我常用的判断方法先用num_workers0跑一遍再用2、4、6、8分别跑看训练速度哪个最快。一般CPU逻辑核心数的一半以下比较稳妥。注意Windows上如果num_workers0却不加启动保护程序直接卡死或报错这个我在第4节专门讲。pin_memory只在使用GPU训练时有用。设为True时DataLoader会把你CPU上的Tensor分配在锁页内存里这样一来GPU拷贝数据的速度会更快。显存够、内存够的前提下建议直接设为True。3.4 collate_fn把零散样本拼成一个批次的关键回调collate_fn是最容易被忽视、却最影响灵活性的参数。DataLoader默认的做法是从__getitem__拿到一批样本把每个位置的元素分别堆成新Tensor。比如batch里有32个(input, label)它会把32个input堆成(32, ...)的Tensor32个label堆成(32, ...)。但现实世界的数据经常不是等长的。比如自然语言处理里一句话有10个词另一句话有50个词你没法直接堆成矩形张量。这时就要自定义collate_fn在拼接之前做一个“补齐”操作def collate_fn(batch): inputs, labels zip(*batch) seq_lens [x.size(0) for x in inputs] max_len max(seq_lens) padded torch.zeros(len(inputs), max_len, inputs[0].size(1)) for i, x in enumerate(inputs): padded[i, :x.size(0)] x return padded, torch.stack(labels)这种补齐逻辑放到底层的Dataset.__getitem__里做是做不到的因为单条样本不知道同batch里其它样本的长度。数据增强里的固定长度裁剪比如统一裁剪成max_len当然也可以但会丢掉长句信息。用collate_fn做动态补齐既保留原始长度又满足了Tensor堆叠要求。更常见的用法是在collate_fn里做数据增强比如把音频加上随机噪声或者把图片做MixUp。它知道你手头这个batch的所有信息适合做batch级别的操作。这块逻辑写到DataLoader里训练循环就特别清爽。4. 跑起来之后才会踩的四个坑我一个一个说代码写出来能跑和跑得稳、跑得快是两回事。我把自己踩过的坑以及身边初学者常遇到的问题集中放到这一节。4.1 Windows下num_workers0会卡死的真正原因如果你在Windows上设num_workers4运行时发现程序卡住不动或者直接报错别怀疑是DataLoader坏了。这个问题的根源在于Windows下多进程是通过“重新导入脚本文件”的方式来启动子进程的。如果你的训练脚本不是写在if __name__ __main__:这个保护块里那么当子进程启动时会重新执行整个脚本相当于又创建了一批子进程然后无限递归下去直到进程数爆炸。解决办法有两条。第一把你的训练逻辑放进main()函数然后在文件末尾写if __name__ __main__: main()第二如果不想把代码包在函数里可以硬编码num_workers0来避开多进程。但这样会损失并行加载的速度所以正规做法还是第一种。这一点很多教程里压根不提但零基础遇到几乎必踩。4.2 worker进程里的随机种子当num_workers0时每个子进程都有自己的随机数生成器。如果不做处理子进程之间的随机状态可能是一样的这会导致一个问题如果你在__getitem__里做了随机数据增强同一个epoch里每个epoch之间是独立的但epoch内部不同batch之间增强结果可能高度重合。训练效果可能没有大问题但如果你要做严格的复现实验这一条很致命。解决办法是在创建DataLoader时传入一个worker_init_fn在每个worker进程启动时重置随机种子def worker_init_fn(worker_id): torch.manual_seed(42 worker_id) loader DataLoader(dataset, num_workers4, worker_init_fnworker_init_fn)这个技巧对于需要多次实验对比的场景特别重要。你可以把不同设置的随机种子固定下来让模型初始化和数据增强完全确定这样每次跑出来的结果差异只来自模型本身而不是数据的随机性。4.3 验证集的数据泄漏这个问题出现得比想象中频繁。很多同学的代码结构是先定义一个Dataset类再实例化训练集和验证集。如果这个Dataset类里带了随机增强比如随机裁剪、翻转而你验证时也直接用了同一个类那么验证集看到的数据也被随机改变了。这看起来没啥实际上会让验证集无法真实反映训练效果。正确的做法是让验证集使用一套不含随机增强的transform。最常见的是在Dataset的__init__里增加一个train或transform参数根据训练/验证状态选择不同的处理方式。像前面写的transform那段训练时用带翻转的组合验证时只用Resize ToTensor Normalize。验证集数据必须固定、可复现否则你无法判断模型是不是真的学到了东西。我还遇到过一种更隐蔽的泄漏collate_fn里做了归一化用的是整个数据集计算出来的均值和标准差。如果这个均值和标准差是在包含验证集的数据上算出来的那验证集的分布就被“偷看”了。严格的做法是用训练集单独计算归一化参数验证集直接用同一套参数但计算过程不涉及验证集的数据。4.4 重置迭代器与epoch的闭环DataLoader本质上是围绕__len__和__getitem__实现的迭代器每次迭代会取一个batch直到取完。取完之后这个DataLoader对象并不会自动回到开头你需要在下一次eopoch训练轮次时重新获取迭代器。写训练循环的时候最常见的一种写法是for epoch in range(num_epochs): for batch in loader: ...这个写法本身没问题因为外层循环每跑一次for batch in loader都会重新创建迭代器。但如果你为了某种目的把DataLoader定义成了内部共享的同一个迭代器比如it iter(loader)然后两个epoch都从这个it里取那么第二个epoch一开始就是空的了循环直接不执行。这个坑很隐蔽因为代码不报错但训练到一半loss不动了。我的习惯是每次epoch都把loader重新交给enumerate()或者iter()不要让迭代器跨epoch复用。还有一种情况是自定义Sampler实现样本不均衡时的采样这时需要特别关注迭代器状态和DataLoader配合起来更复杂零基础可以先绕开等理解了基础流程再来接触。5. 常用Containers那些被忽略却每天都在用的容器模块数据管线打通了模型结构这一层也不能含糊。PyTorch里有一组叫“容器”的模块很多人一开始根本不知道它们的存在直到看到别人的代码里出现nn.Sequential、nn.ModuleList才开始问这是什么。容器的核心作用是帮你把若干子模块组织起来统一管理参数、统一迁移设备和保存加载。没有容器的时候你只能手动维护一个成员变量列表但这样PyTorch无法识别你的模型里有哪些参数。容器解决了这个问题——只要子模块被放进容器它们就自动登记进模型的parameters()和state_dict()里。5.1 nn.Sequential给单一“串行网络”的直通管道nn.Sequential是最常用的容器。它按顺序存放一组层forward时自动按顺序执行。适合那种一条直线走到底的网络比如一个简单的全连接网络import torch.nn as nn class MLP(nn.Module): def __init__(self, input_dim64, hidden_dim128, output_dim10): super().__init__() self.net nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(inplaceTrue), nn.Linear(hidden_dim, output_dim), ) def forward(self, x): return self.net(x)这个写法的好处是代码紧凑层的顺序一目了然。当你需要插入一层时在括号里加一行就行不用单独写forward逻辑。但它的局限也在这里只能表达严格的串行计算。如果中间有分支、有跳跃连接比如ResNet那种shortcutSequential就套不住了需要自定义forward。所以我的经验是结构简单的模型用Sequential结构复杂了一定要自己写forward别硬套。5.2 nn.ModuleList为什么不能直接用Python的list很多初学者想当然地写self.layers [nn.Linear(64, 128) for _ in range(5)]然后训练时发现模型可以跑loss也在下降但保存参数时发现state_dict里根本没有这5个Linear层的权重。原因很简单Python的原生list不是一个Module容器PyTorch不会去遍历它找出你存进去的层。于是这些层成了“游离”的模块——forward照用但参数不在模型管理范围内。正确的做法是换成nn.ModuleListhidden_sizes [64, 128, 256, 128, 32] self.layers nn.ModuleList() for i in range(len(hidden_sizes) - 1): self.layers.append(nn.Linear(hidden_sizes[i], hidden_sizes[i 1])) def forward(self, x): for layer in self.layers: x torch.relu(layer(x)) return xnn.ModuleList和原生list用法几乎一样索引、追加、迭代都支持。唯一的区别是里面的每个子模块会被自动注册到当前Module上成为参数集合的一部分。因此调用model.to(device)、model.state_dict()、model.parameters()时它们都会被正确处理。在什么场景下一定要用ModuleList当你有“N个相同结构的块”时比如Transformer里有12个编码层图像网络里有多个Stage手动把每层写出来会又长又容易错用for循环生成层放进ModuleList既干净又省事。5.3 ModuleDict与ParameterList的适配场景nn.ModuleDict是按名字存模块的容器适合那种需要“按分支名取模块”的网络。比如一个多模态模型不同输入走不同的分支可以这样组织self.branch nn.ModuleDict({ image: nn.Sequential(nn.Conv2d(3, 64, 3), nn.ReLU(inplaceTrue)), text: nn.Sequential(nn.Linear(100, 64), nn.ReLU(inplaceTrue)), audio: nn.Sequential(nn.Linear(80, 64), nn.ReLU(inplaceTrue)), }) # forward里根据key取分支 out self.branch[input_type](x)用原生字典去存就会出现和原生list一样的问题字典里的模块不会自动注册训练时参数不会被优化。换成nn.ModuleDict问题就消失了。还有一种容器nn.ParameterList、nn.ParameterDict存放的是Parameter本身而不是Module。如果你的网络里有一些不是标准层的可学习参数比如多任务学习里每个任务一个独立的线性映射矩阵就可以用ParameterList来管理。它和ModuleList的区别是这里存的是要背梯度的参数张量不是完整的层。self.scale_params nn.ParameterList([ nn.Parameter(torch.ones(1)) for _ in range(num_tasks) ])Parameters放进ParameterList之后模型在优化器里能自动把参数找出来不会因为没被某个Module直接持有而漏掉。6. 一个同时用上DataLoader和Containers的完整流程示例讲到这里数据加载和容器其实已经打通了思路。我搭一个简单的多任务回归例子把两部分一起串一遍你看着这段代码会感觉比单独看理论强很多。任务设定输入是64维向量输出是两个维度的预测值。数据集照旧随便造。模型由一个共享主干和两个任务分支组成主干用Sequential分支用ModuleDict最后用DataLoader做批量训练。import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader class PairDataset(Dataset): def __init__(self, num_samples20000, input_dim64): self.x torch.randn(num_samples, input_dim) self.y1 (self.x[:, 0] self.x[:, 1]).unsqueeze(1) self.y2 (self.x[:, 2] * self.x[:, 3] 1).unsqueeze(1) def __len__(self): return len(self.x) def __getitem__(self, idx): return self.x[idx], self.y1[idx], self.y2[idx] class MultiTaskNet(nn.Module): def __init__(self, input_dim64): super().__init__() self.feature nn.Sequential( nn.Linear(input_dim, 128), nn.ReLU(inplaceTrue), nn.Linear(128, 32), nn.ReLU(inplaceTrue), ) self.head nn.ModuleDict({ task1: nn.Linear(32, 1), task2: nn.Linear(32, 1), }) def forward(self, x, task_nametask1): feat self.feature(x) return self.head[task_name](feat)训练循环dataset PairDataset() loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers0) model MultiTaskNet() optimizer torch.optim.Adam(model.parameters(), lr1e-3) loss_fn nn.MSELoss() for epoch in range(5): for x, y1, y2 in loader: optimizer.zero_grad() out1 model(x, task1) loss1 loss_fn(out1, y1) out2 model(x, task2) loss2 loss_fn(out2, y2) loss loss1 loss2 loss.backward() optimizer.step() print(fepoch {epoch1}, loss: {loss.item():.4f})这段代码里能看出好几个设计点。Dataset返回三个元素DataLoader默认就会把它们堆成三个批量Tensor。ModuleDict让分支选择变成传字符串结构扩展时只要在字典里加一层forward不用改。整个模型的参数不管放在Sequential里还是ModuleDict里都被model.parameters()一并拿到优化器更新时不会漏掉任何一个子层。如果需要一个严格对比实验可以把DataLoader加上worker_init_fn并且在一个设置了全局随机种子的脚本里运行保证数据顺序和模型初始化都可复现。说完代码我再把这篇的实际体会提炼一下。DataLoader和Containers看起来是两个不相干的主题但它们背后是同一个思想PyTorch里的模块化管理。数据侧Dataset定义“单条数据怎么取”DataLoader定义“整批数据怎么组织”模型侧nn.Module定义“模型内部怎么计算”容器定义“子模块怎么被管理”。把这两条管线理清楚了后面学习迁移学习、分布式训练、自定义模型结构时就不会觉得这些东西是突然冒出来的。我建议你打开一个脚本把我给的这两个片段分别敲一遍再试着改改batch_size、加一个分支、换一个自定义collate_fn。代码跑通之后这一篇才算真正学到手了。
解法)
