谈到PyTorch里的数据加载我见过太多“模型抄得飞快训练一跑就翻车”的情况——而翻车点往往不在网络结构本身而在Dataset和DataLoader这一层。构建数据集这件事看似只是“把数据喂给模型”实际上一旦涉及分类标签映射、多样本维度、多进程加载、内存占用这些细节坑一个比一个深。这篇博文就是要把Dataset和DataLoader这两个组件彻底拆开它们各自负责什么、手写Dataset时该注意哪些门道、DataLoader里每个参数怎么取舍、怎么搭一套能直接用于训练的完整数据管道以及高频报错怎么排查。先说清楚这篇适合谁读。如果你已经能写出一个简单的分类模型但一碰数据准备就心里发虚——不知道__getitem__到底要返回什么、不知道num_workers为什么一开就卡死、不知道变长数据怎么凑成一个batch那么这篇文章是为你准备的。内容完全基于常见的真实项目场景展开代码可以直接复制改造所有关键取舍我都会说明背后的理由而不是让你死记参数表。1. Dataset和DataLoader的分工先搞懂谁干什么活1.1 一个借书的比喻把两件事分开把整个数据准备机制想象成一个巨型图书馆。Dataset就是那个藏书库它只需要做到两件事知道自己馆里一共有多少本书__len__以及当你说“我要第i本书”时能把那一本准确拿出来递给你__getitem__。至于书是按什么分类放的、书页是纸质还是电子屏书库内部怎么折腾都不关外面的事。DataLoader则是馆外的借阅服务台。它不关心书库内部怎么存书只负责按你的要求组织“取书”流程一次取几本batch_size、要不要打乱顺序shuffle、要不要派好几个管理员同时进库去取num_workers、如果最后一轮书不够一组怎么处理drop_last。训练代码每次去找借阅台要数据就能稳定拿到一批形状规整、顺序符合预期的样本。这个比喻能帮你立刻理解一个核心点Dataset定义的是“怎么拿一条数据”DataLoader定义的是“怎么组织一堆数据”。两者的接口高度解耦所以照片在文件夹里、文本在数据库里、特征在CSV里都不影响训练代码里for images, labels in loader这个循环的写法。1.2 为什么非要用这套机制而不是直接读文件很多人第一次接触PyTorch都会有一个疑问直接把所有数据读成一个Tensor扔给模型不行吗小数据集当然可以这么干但往大了走必然崩溃。第一内存瓶颈。把几万张高清图片一次性读进内存机器几乎必挂更别提还要把它们转成Tensor、做归一化、执行数据增强。Dataset的惰性加载机制让每一次__getitem__只处理一条样本用完即走内存占用被压到极低。第二并行潜力。DataLoader开多个num_workers后会有多个子进程同时调用Dataset的__getitem__相当于多个管理员同时在不同书架上取书。读取图片、做数据增强这类CPU密集型操作可以跟GPU上的前向反向计算重叠起来让训练效率提升一个量级。第三扩展性。无论数据源是本地文件夹、云端存储、数据库还是其他自定义格式只要实现好Dataset那三个接口上层代码就不用变。这也是框架设计的高明之处把易变的部分数据来源和预处理逻辑封装在子类里把稳定的部分组batch、打乱、并行留在框架里。2. 手写第一个Dataset类核心模板与三种实战写法2.1 必须实现的三个魔法方法自定义Dataset必须继承torch.utils.data.Dataset并且实现三个方法缺一不可__init__(self)初始化阶段。通常在这里准备文件路径列表、标签列表、transform对象等。注意这里做的是“轻量准备”不是把全部数据load进内存。__len__(self)返回数据集的样本总数。DataLoader靠它计算一个epoch有多少个batch。__getitem__(self, idx)接收一个索引返回一条样本。样本可以是(图像, 标签)元组也可以是包含更多信息的字典。其中__getitem__是真正干重活的函数也是绝大多数新手容易写乱的地方。一个稳定的__getitem__设计应该做到“只依赖idx这一个输入不依赖任何可变全局状态”否则在多进程加载时很容易出现数据错乱。2.2 图像分类Dataset实战最经典的一种先给一个通用的图像分类Dataset模板。它针对的是最常见的目录结构根目录下有若干类别文件夹每个文件夹里是对应类别的图片。import os from PIL import Image from torch.utils.data import Dataset class ImageFolderDataset(Dataset): 从标准文件夹结构构建图像分类数据集。 目录结构示例 root_dir/ cat/001.jpg cat/002.jpg dog/001.jpg def __init__(self, root_dir, transformNone): super().__init__() self.root_dir root_dir self.transform transform # 类别顺序必须先固定第二次运行结果才一致 self.classes sorted(os.listdir(root_dir)) self.class_to_idx {cls_name: i for i, cls_name in enumerate(self.classes)} self.samples [] for cls_name in self.classes: cls_dir os.path.join(root_dir, cls_name) if not os.path.isdir(cls_dir): continue for file_name in os.listdir(cls_dir): if file_name.lower().endswith((.png, .jpg, .jpeg)): self.samples.append( (os.path.join(cls_dir, file_name), self.class_to_idx[cls_name]) ) def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] # 统一转RGB避免混入灰度图或RGBA四通道图引发维度错乱 image Image.open(img_path).convert(RGB) if self.transform is not None: image self.transform(image) return image, label这段代码里有几个细节是平时文档里不太会强调的。os.listdir在Windows和Linux上返回的目录顺序并不保证一致所以我先用sorted把类别名排好序再建立索引映射。不要小看这一步一旦class_to_idx的映射顺序在不同机器上不一致就会造成训练和评估时标签完全对不上。每次__getitem__都去os.listdir遍历目录是不现实的所以我在__init__里把(路径, 标签)列表一次性建好这属于“轻量准备”维护一个元组列表的开销极小却能让__getitem__变成纯O(1)的查表加读取操作。Image.open本身是惰性的真正读像素是在convert(RGB)或transform执行时才发生所以不用担心__init__阶段就把图片读进来了。2.3 表格回归Dataset实战数值型数据的另一种写法图像数据量大、读取慢但表格数据通常是轻量的。处理表格数据时直接在__init__里把整个CSV读进去反而是效率最高的方案。import numpy as np import pandas as pd import torch from torch.utils.data import Dataset class TabularDataset(Dataset): 适用于表格型回归/分类任务的数据集。 def __init__(self, csv_path, feature_cols, target_col): super().__init__() df pd.read_csv(csv_path) # 转成float32避免模型计算时dtype不一致 self.features df[feature_cols].values.astype(np.float32) self.targets df[target_col].values.astype(np.float32) assert len(self.features) len(self.targets), \ feature和target长度不一致请检查数据 def __len__(self): return len(self.features) def __getitem__(self, idx): x torch.from_numpy(self.features[idx]) y torch.tensor(self.targets[idx]) return x, y这里特意用了NumPy数组而不是DataFrame切片是因为df[col].iloc[idx]在每次__getitem__调用时都走一遍pandas的索引逻辑慢且容易产生额外开销。转成NumPy后直接整数索引切片性能会明显好一截。至于表格数据要不要像图像那样惰性读取我的建议是总量在几个GB以内直接全读进内存。表格特征一般就几十列一条样本撑死几百字节全量加载对内存的压力可以忽略。反过来如果每条样本包含文本、向量甚至大块JSON那就需要重新评估别无脑套模板。2.4 一个新手最容易踩的坑在init阶段把所有数据塞进内存这是个很典型的行为——新手为了图省事在__init__里把所有图片转成numpy数组放进一个list。小项目demo跑得飞快数据一涨就直接内存崩溃。原因在于DataLoader一旦开多个num_workers子进程会持有Dataset对象的引用数据量越大内存压力放大得越明显。正确思路是__init__里只维护“怎么找到数据”的索引比如路径列表、行号列表__getitem__里才真正执行“读取数据 预处理”。把重IO放getitem还能顺带享受多进程并行加载的红利——Every worker都在并行做读图和图像增强等于把预处理这步也并行化了。有一种例外数据总量较小且预处理开销大可以考虑在第一次遍历时把结果缓存到内存第二次开始直接命中缓存。这是空间换时间的权衡跟无脑塞内存有本质区别后面进阶部分会展开讲。3. DataLoader参数逐个拆batch_size只是开胃菜3.1 核心参数配置速查表DataLoader的可调参数不少但真正在日常项目中频繁用到的最核心就这几个。先把它们的作用和默认值列出来再逐个展开。参数默认值作用日常建议batch_size1每批样本数量按显存和模型设计来调整shuffleFalse是否在每个epoch开始前打乱顺序训练集True验证/测试集Falsenum_workers0加载数据的子进程数调0试跑通正式训练往4-8调collate_fnNone把多条样本拼成一个batch的策略变长数据必须自定义drop_lastFalse最后一个不满batch_size的batch是否丢弃训练集建议Truepin_memoryFalse是否使用锁页内存GPU训练建议Trueprefetch_factor2每个worker预先取几批数据一般默认即可samplerNone自定义采样顺序和shuffle互斥特殊采样需求时用3.2 shuffle、sampler和batch_sampler怎么选shuffleTrue的本质是在每个epoch开始前对整个数据集的索引重新洗牌。它的优点是不需要你关心任何采样细节框架自动处理。但注意shuffle和sampler不能同时设置因为shuffle本身就是一种特定sampler。什么时候需要自定义sampler类别极度不均衡时。假设一个分类任务里样本A有1万条样本B只有100条普通shuffle会导致每个batch里几乎全是A模型学不到B的特征。这时候用WeightedRandomSampler按类别权重采样让每个类别被抽到的概率更接近。from torch.utils.data import WeightedRandomSampler # 每个类别的权重 1 / 该类别样本数样本越少的类别权重越大 class_counts {} for _, label in train_dataset: class_counts[label] class_counts.get(label, 0) 1 sample_weights [1.0 / class_counts[label] for _, label in train_dataset] sampler WeightedRandomSampler( sample_weights, num_sampleslen(train_dataset), replacementTrue ) train_loader DataLoader( train_dataset, batch_size32, samplersampler, # 注意不要再设置shuffleTrue num_workers4, )batch_sampler更进一步直接控制每个batch包含哪些索引。它的优先级最高但同时指定了它batch_size和shuffle都会失效。实际业务中用得少但当你需要实现“同一个batch内样本类别分布固定”之类的要求时就得靠它。3.3 num_workers到底设多少实测下来什么感受num_workers是新手最容易“调爆”的参数。很多人一上来就写num_workers8然后发现要么程序卡死要么跑得比num_workers0还慢要么直接报worker崩溃。原因很简单多进程不是免费的进程启动有开销进程间数据传递有序列化消耗。我建议按这个节奏来调先用num_workers0跑通全流程确认数据和模型逻辑没问题。图像数据规模中等几千张开num_workers4通常够了。数据量大、transform复杂逐步往上加到8、12、16同时看CPU占用率和每个epoch耗时。当CPU占用接近100%且训练速度不再提升说明已经到了瓶颈继续加是没有意义的。有个容易忽略的点Windows系统下开启多进程必须有if __name__ __main__:保护否则子进程会递归执行主模块里的代码导致崩溃。很多人在IDE里直接跑脚本碰到的worker崩溃八成是这个问题。还有一个硬件层面的坑如果数据放在机械硬盘上随机读小图非常慢这时候盲目加大num_workers反而会让磁盘IO变成新的瓶颈。比起加进程把数据放到SSD或者换一种更顺序化的读取方式效果立竿见影。3.4 pin_memory和drop_last这对低调参数pin_memoryTrue会让DataLoader把数据放到锁页内存里这个操作本身不影响CPU上跑的数据内容但当你执行tensor.to(device, non_blockingTrue)时CPU到GPU的拷贝可以走异步路径跟当前线程的其他计算重叠从而稍微降低数据传输的阻塞时间。代价是锁页内存管理比普通内存更严格占用过多会拖慢系统。drop_lastTrue则负责丢掉最后一组不够batch_size的样本。为什么训练时建议开因为很多框架的BatchNorm等层会依赖稳定的batch统计量如果最后一个batch只有1条样本这种batch的统计量根本没有参考意义还可能让loss曲线出现异常抖动。验证集则无所谓保持False让所有样本都能被评估到。4. 从零搭一套完整的图像分类数据管道4.1 文件夹结构设计与transform配置动手写代码前先把目录结构设计好。下面这个结构是最常见、也最方便我上面那个Dataset直接使用的。image_classification/ data/ train/ cat/001.jpg cat/002.jpg dog/001.jpg val/ cat/001.jpg dog/001.jpg train.py然后是transform的配置。这里有一个关键认知训练集的transform应当以“增强”为主验证集的transform应当以“固定”为主。训练时做随机裁剪、随机翻转让模型看到更多变体验证时统一resize到固定尺寸再中心裁剪保证每次评估的是同一个标准下的结果。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.3), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])注意ToTensor的位置一定要在所有像素级增强之后、Normalize之前。因为ToTensor会把像素值从0-255缩放到0-1区间并变换通道顺序为CHW而Normalize则在此基础上按通道做标准化。如果顺序反了Normalize的均值和方差就会作用绝对错误的像素范围上。这组mean和std来自ImageNet预训练统计。使用ImageNet预训练模型时输入最好和它训练时的分布保持一致所以默认沿用这套值。如果你是自己从零训练模型可以统计自己数据的均值和方差来替换不过说实话直接沿用预训练统计对多数场景也够用。4.2 完整代码串联Dataset DataLoader 训练循环有了Dataset类和transform接下来把它们拼成完整的训练管道。下面这段代码可以直接作为二分类图像任务的起点模板。import os import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision.models import resnet18 # 沿用前面定义的 ImageFolderDataset # 沿用前面的 train_transform / val_transform def build_dataloaders(root_dir, batch_size64, num_workers4): train_ds ImageFolderDataset( os.path.join(root_dir, train), transformtrain_transform ) val_ds ImageFolderDataset( os.path.join(root_dir, val), transformval_transform ) train_loader DataLoader( train_ds, batch_sizebatch_size, shuffleTrue, num_workersnum_workers, drop_lastTrue, pin_memoryTrue, ) val_loader DataLoader( val_ds, batch_sizebatch_size, shuffleFalse, num_workersnum_workers, pin_memoryTrue, ) return train_loader, val_loader def main(): device torch.device(cuda if torch.cuda.is_available() else cpu) train_loader, val_loader build_dataloaders(data, batch_size64) model resnet18(num_classes2).to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) for epoch in range(10): model.train() for images, labels in train_loader: images images.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() # 验证阶段 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images images.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) outputs model(images) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) print(fepoch {epoch 1}, val acc: {correct / total:.4f}) if __name__ __main__: main()完整代码里有一条很关键的主线dataset负责产出“一条样本”dataloader负责把它变成“一个batch”训练循环只认“一个batch”。将来不管改成文本任务、语音任务还是多模态任务这一层的衔接方式不会变。4.3 验证数据管道的正确性别等训练跑完才发现数据有问题数据管道写得对不对不应该等到训练完看acc才去猜。正确做法是在正式训练前花两分钟检查一下输出的形状、类型和数值范围。def inspect_loader(loader, name): images, labels next(iter(loader)) print(f{name}:) print(f images shape: {images.shape}) print(f labels shape: {labels.shape}) print(f images dtype: {images.dtype}) print(f label dtype: {labels.dtype}) print(f image value range: [{images.min():.3f}, {images.max():.3f}]) inspect_loader(train_loader, train) inspect_loader(val_loader, val)一个正常的输出大致长这样train: images shape: torch.Size([64, 3, 224, 224]) labels shape: torch.Size([64]) images dtype: torch.float32 label dtype: torch.int64 image value range: [-2.518, 2.645]看到[64, 3, 224, 224]就说明batch、通道、高度、宽度四个维度都对dtype是float32不会引发模型里的类型错误图像值出现负数说明Normalize生效了。如果值还在0到1之间说明你很可能忘了把Normalize接在transform列表里。再进一步可以用torchvision.utils.save_image把增强后的图片保存下来肉眼看一看增强效果有没有过度。数据增强是把双刃剑强度拉太满会让图中物体变形到无法辨认视觉效果比线性指标敏感得多。from torchvision.utils import save_image images, labels next(iter(train_loader)) save_image(images[:8], aug_vis.png, nrow4, normalizeTrue, value_range(-3, 3))打开aug_vis.png如果里面图片内容都认不出来多半是ColorJitter的幅度设大了或者RandomResizedCrop裁剪比例过狠需要回调。4.4 磁盘类型对数据管道的影响有人在这上面被坑惨了同样的代码、同样的num_workers在机械硬盘和固态硬盘上的表现可能相差数倍。机械硬盘的随机读取能力很差而图像数据集恰恰是大量小文件的随机读取这在机械硬盘上是灾难性的。如果数据量不算特别大一个实操有效的优化是把图片打包成一个大文件存储或者干脆在第一次遍历时缓存到内存。内存的读取速度比任何盘都快几个数量级只要内存放得下就把缓存空间利用起来。如果内存放不下再退而求其次把预处理后的特征事先算好存成本地.pt文件训练时只读特征文件。这样做训练阶段直接绕开原图和transform速度提升非常明显。代价是需要额外跑一遍离线预处理且磁盘占用会变大。对“样本很大、重复训练很多轮”的场景这笔投入很划算。5. 数据加载常见问题与排查实录5.1 高频问题速查表数据管道的坑看起来千奇百怪实际上高度重复。把最常见的情况整理成一张表遇到问题先对号入座。现象常见原因排查方向每个epoch拿到的第一个batch完全一样shuffleFalse或使用了sampler但没有手动打乱确认训练集shuffleTrue不要同时指定sampler训练时程序卡死无报错Windows下没加if __name__ __main__保护把创建数据集和loader的代码全部放进main()里报错DataLoader worker exited unexpectedly子进程在getitem里抛异常主进程看不到完整堆栈先设num_workers0跑出完整traceback报错stack expects each tensor to be equal size样本维度不一致比如图像有灰图也有RGB图检查convert(RGB)是否加了变长数据写自定义collate_fn报错Expected all tensors to be on the same device数据或模型有一侧忘了.to(device)打印images.device和next(model.parameters()).device内存持续上涨直到崩溃__getitem__里缓存了太多数据或初始化时全量加载审视Dataset里是否有大的list/cache确认没有全局引用GPU利用率忽高忽低num_workers不够、数据在机械硬盘、transform太重看CPU和磁盘IO占用再针对性调整验证集所有标签预测同一个类别可能是class_to_idx建错或数据顺序没打乱让模型学到捷径打印训练前labels的取值分布和class_to_idx对照5.2 两个必查的调试点一个batch和可视化的增强结果排查数据管道问题最有效的第一招永远是“打印一个batch”。next(iter(loader))这个操作成本极低却能让你立刻看清形状、dtype、数值范围这些最基础但最容易出错的信息。之前有个案例某人模型训练acc死活上不去折腾了很久结果一查是__getitem__返回的标签没转成int64CrossEntropyLoss计算时类型不匹配。这种问题在打印labels的一瞬间就能发现。第二个必查点是增强后的可视化。save_image保存几张增强后的图片肉眼看一遍很多异常原因会比看代码更快暴露。比如RandomResizedCrop的scale参数如果设置不当可能让图片被裁到只剩一个角ColorJitter参数过大会让整个色调失真。这些“增强过度”的问题不通过可视化很难从数值上感知。排查数据管道时永远先怀疑数据再怀疑模型。模型结构可以用预训练权重兜底但错误的数据喂进去再好的模型也只会越训越偏。5.3 排查思路先把所有条件降到最保守档当遇到一个相当诡异、一时摸不着头脑的问题时我的习惯是先把所有可能引入干扰的配置全部去掉。具体操作是num_workers0只保留单进程加载。pin_memoryFalse跳过锁页内存。shuffleFalse固定数据顺序。batch_size调小比如8。transform暂时只保留ToTensor()先不做任何数据增强。这样配置跑一轮如果能跑通说明逻辑本身没问题问题出在并行或增强配置上如果连这样都报错那问题大概率就在Dataset的代码逻辑本身。然后一条一条把配置加回去每次只改一个变量就能精准定位是哪一个环节引入了问题。这个排查法虽然朴素但比同时改三个配置然后“瞎猜”要高效得多。还有一种变长数据场景需要专门提一下因为很多跑NLP或目标检测任务的人会碰到。默认的collate_fn会把多个样本直接torch.stack但不同样本长度不一致时这个stack必然失败。解决办法是自定义collate_fn手动把样本pad到当前batch内的最大长度然后把真实的attention_mask或原始长度一起返回。def collate_fn(batch): input_ids [item[input_ids] for item in batch] attention_mask [item[attention_mask] for item in batch] labels torch.stack([item[label] for item in batch]) max_len max(len(ids) for ids in input_ids) padded_ids torch.zeros(len(batch), max_len, dtypetorch.long) padded_mask torch.zeros(len(batch), max_len, dtypetorch.long) for i, ids in enumerate(input_ids): length len(ids) padded_ids[i, :length] ids padded_mask[i, :length] attention_mask[i] return { input_ids: padded_ids, attention_mask: padded_mask, label: labels, }这个函数的核心思想是先把每条样本的input_ids和attention_mask对齐到同一个长度空间再返回一个字典供模型使用。用到collate_fn的项目十有八九是文本、目标框这类“每条样本长度天然不一致”的数据提前把这套代码准备着不吃亏。6. 进阶用法让数据管道再快一点、再稳一点6.1 用空间换时间内存缓存与特征预计算一个训练跑好多轮的项目数据管道优化能带来非常直观的收益。最省事的优化是加内存缓存第一次访问某个索引时把该样本处理完存起来后续再访问直接命中缓存。class CachedDataset(Dataset): 给任意Dataset加一层缓存包装适合数据总量不大、预处理费时的场景。 def __init__(self, dataset): super().__init__() self.dataset dataset self.cache {} def __len__(self): return len(self.dataset) def __getitem__(self, idx): if idx not in self.cache: self.cache[idx] self.dataset[idx] return self.cache[idx]用的时候直接包一层train_ds CachedDataset(ImageFolderDataset(...))。注意这个方案只适合“样本小、总量放得进内存”的情况。如果一条样本就是一张高清大图加一堆标注缓存会瞬间吃光内存那就得换特征预计算方案——提前把不做随机增强的预处理结果比如模型倒数第二层特征离线算好存成文件训练时直接读特征。这种“离线预处理 在线只读特征”的思路本质上是把重复计算一次性做完。对多轮训练来说绝对划算。我实际测试过某些图像场景下这一招能把数据管道的时间从占训练总耗时的40%压到5%以内。6.2 分布式训练的DataLoader注意点如果你的训练脚本涉及多卡或多机并行DataLoader的普通shuffle就不够用了。分布式训练需要每个进程拿到数据的不同部分同时还要保证所有进程合在一起能覆盖完整的数据集。这个工作由DistributedSampler完成而它也带来了一个额外要求每个epoch必须调用set_epoch否则每个epoch的打乱顺序会完全一样。from torch.utils.data.distributed import DistributedSampler sampler DistributedSampler(train_dataset, shuffleTrue) train_loader DataLoader( train_dataset, batch_size64, samplersampler, num_workers4, ) for epoch in range(10): train_loader.sampler.set_epoch(epoch) for images, labels in train_loader: # 训练逻辑 pass这个set_epoch非常容易被漏掉。一旦漏掉每个epoch的数据划分虽然错开了但打乱后重新划分的随机顺序不随epoch变化模型在每个epoch都会按同样的顺序看到数据等于变相削弱了随机化带来的泛化收益。6.3 进度条不只是好看它帮你发现数据管道瓶颈训练时配一个进度条除了肉眼舒服还有一个实用功能能直观看到每个epoch的耗时变化以及当前每秒处理多少个batch。当数据管道是瓶颈时进度条的每秒batch数会明显异常。from tqdm import tqdm for epoch in range(num_epochs): pbar tqdm(train_loader, descfEpoch {epoch 1}) for images, labels in pbar: images images.to(device, non_blockingTrue) labels labels.to(device, non_blockingTrue) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() pbar.set_postfix(lossloss.item())如果GPU利用率偏低但进度条推得很慢瓶颈多半在数据加载侧如果GPU利用率很高但进度条依然不快说明模型计算本身就重数据管道已经喂得够快了。这个“看进度条判断瓶颈在哪端”的小技巧我在实际项目里判断数据管道优化空间时几乎是第一直觉手段。6.4 数据管道是地基花时间打磨永远值得再分享一个我已经固化成习惯的流程。每接手一个新的数据集我不会直接开足马力去调大num_workers或追求高batch_size而是先写一个极简的inspect脚本打印前几个batch的shape、dtype、数值范围保存两张增强后的图片看一眼结果。等确认一切合理后再进入正式训练循环。这个小习惯帮我避开了至少两位数的“训练跑完才发现数据有问题”的尴尬场景。如果你此刻正在一个新的PyTorch项目上起步我建议你也把这个inspect环节固定下来。别急着把模型写多复杂先把 Dataset 和 DataLoader 这一层打磨扎实。数据管道是深度学习中所有结论的地基地基稳了后面换模型、调损失函数、做实验对比才不白折腾。
