谈到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, transform=None): 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怎么选shuffle=True的本质，是在每个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_samples=len(train_dataset), replacement=True ) train_loader = DataLoader( train_dataset, batch_size=32, sampler=sampler, # 注意不要再设置shuffle=True num_workers=4, )batch_sampler更进一步，直接控制每个batch包含哪些索引。它的优先级最高，但同时指定了它，batch_size和shuffle都会失效。实际业务中用得少，但当你需要实现“同一个batch内样本类别分布固定”之类的要求时，就得靠它。3.3 num_workers到底设多少，实测下来什么感受num_workers是新手最容易“调爆”的参数。很多人一上来就写num_workers=8，然后发现：要么程序卡死，要么跑得比num_workers=0还慢，要么直接报worker崩溃。原因很简单：多进程不是免费的，进程启动有开销，进程间数据传递有序列化消耗。我建议按这个节奏来调：先用num_workers=0跑通全流程，确认数据和模型逻辑没问题。图像数据规模中等（几千张），开num_workers=4通常够了。数据量大、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_memory=True会让DataLoader把数据放到锁页内存里，这个操作本身不影响CPU上跑的数据内容，但当你执行tensor.to(device, non_blocking=True)时，CPU到GPU的拷贝可以走异步路径，跟当前线程的其他计算重叠，从而稍微降低数据传输的阻塞时间。代价是锁页内存管理比普通内存更严格，占用过多会拖慢系统。drop_last=True则负责丢掉最后一组不够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(brightness=0.3, contrast=0.3, saturation=0.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_size=64, num_workers=4): train_ds = ImageFolderDataset( os.path.join(root_dir, 'train'), transform=train_transform ) val_ds = ImageFolderDataset( os.path.join(root_dir, 'val'), transform=val_transform ) train_loader = DataLoader( train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True, pin_memory=True, ) val_loader = DataLoader( val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True, ) 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_size=64) model = resnet18(num_classes=2).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) for epoch in range(10): model.train() for images, labels in train_loader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) 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_blocking=True) labels = labels.to(device, non_blocking=True) outputs = model(images) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) print(f"epoch {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', nrow=4, normalize=True, value_range=(-3, 3))打开aug_vis.png，如果里面图片内容都认不出来，多半是ColorJitter的幅度设大了或者RandomResizedCrop裁剪比例过狠，需要回调。4.4 磁盘类型对数据管道的影响：有人在这上面被坑惨了同样的代码、同样的num_workers，在机械硬盘和固态硬盘上的表现可能相差数倍。机械硬盘的随机读取能力很差，而图像数据集恰恰是大量小文件的随机读取，这在机械硬盘上是灾难性的。如果数据量不算特别大，一个实操有效的优化是：把图片打包成一个大文件存储，或者干脆在第一次遍历时缓存到内存。内存的读取速度比任何盘都快几个数量级，只要内存放得下，就把缓存空间利用起来。如果内存放不下，再退而求其次，把预处理后的特征事先算好存成本地.pt文件，训练时只读特征文件。这样做训练阶段直接绕开原图和transform，速度提升非常明显。代价是需要额外跑一遍离线预处理，且磁盘占用会变大。对“样本很大、重复训练很多轮”的场景，这笔投入很划算。5. 数据加载常见问题与排查实录5.1 高频问题速查表数据管道的坑看起来千奇百怪，实际上高度重复。把最常见的情况整理成一张表，遇到问题先对号入座。现象常见原因排查方向每个epoch拿到的第一个batch完全一样shuffle=False，或使用了sampler但没有手动打乱确认训练集shuffle=True，不要同时指定sampler训练时程序卡死无报错Windows下没加if __name__ == '__main__'保护把创建数据集和loader的代码全部放进main()里报错DataLoader worker exited unexpectedly子进程在getitem里抛异常，主进程看不到完整堆栈先设num_workers=0跑出完整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__返回的标签没转成int64，CrossEntropyLoss计算时类型不匹配。这种问题在打印labels的一瞬间就能发现。第二个必查点是增强后的可视化。save_image保存几张增强后的图片，肉眼看一遍，很多异常原因会比看代码更快暴露。比如RandomResizedCrop的scale参数如果设置不当，可能让图片被裁到只剩一个角；ColorJitter参数过大会让整个色调失真。这些“增强过度”的问题，不通过可视化很难从数值上感知。排查数据管道时，永远先怀疑数据，再怀疑模型。模型结构可以用预训练权重兜底，但错误的数据喂进去，再好的模型也只会越训越偏。5.3 排查思路：先把所有条件降到最保守档当遇到一个相当诡异、一时摸不着头脑的问题时，我的习惯是先把所有可能引入干扰的配置全部去掉。具体操作是：num_workers=0，只保留单进程加载。pin_memory=False，跳过锁页内存。shuffle=False，固定数据顺序。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, dtype=torch.long) padded_mask = torch.zeros(len(batch), max_len, dtype=torch.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, shuffle=True) train_loader = DataLoader( train_dataset, batch_size=64, sampler=sampler, num_workers=4, ) 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, desc=f'Epoch {epoch + 1}') for images, labels in pbar: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() pbar.set_postfix(loss=loss.item())如果GPU利用率偏低但进度条推得很慢，瓶颈多半在数据加载侧；如果GPU利用率很高但进度条依然不快，说明模型计算本身就重，数据管道已经喂得够快了。这个“看进度条判断瓶颈在哪端”的小技巧，我在实际项目里判断数据管道优化空间时，几乎是第一直觉手段。6.4 数据管道是地基，花时间打磨永远值得再分享一个我已经固化成习惯的流程。每接手一个新的数据集，我不会直接开足马力去调大num_workers或追求高batch_size，而是先写一个极简的inspect脚本，打印前几个batch的shape、dtype、数值范围，保存两张增强后的图片看一眼结果。等确认一切合理后，再进入正式训练循环。这个小习惯帮我避开了至少两位数的“训练跑完才发现数据有问题”的尴尬场景。如果你此刻正在一个新的PyTorch项目上起步，我建议你也把这个inspect环节固定下来。别急着把模型写多复杂，先把 Dataset 和 DataLoader 这一层打磨扎实。数据管道是深度学习中所有结论的地基，地基稳了，后面换模型、调损失函数、做实验对比，才不白折腾。
谈到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, transform=None): 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怎么选shuffle=True的本质，是在每个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_samples=len(train_dataset), replacement=True ) train_loader = DataLoader( train_dataset, batch_size=32, sampler=sampler, # 注意不要再设置shuffle=True num_workers=4, )batch_sampler更进一步，直接控制每个batch包含哪些索引。它的优先级最高，但同时指定了它，batch_size和shuffle都会失效。实际业务中用得少，但当你需要实现“同一个batch内样本类别分布固定”之类的要求时，就得靠它。3.3 num_workers到底设多少，实测下来什么感受num_workers是新手最容易“调爆”的参数。很多人一上来就写num_workers=8，然后发现：要么程序卡死，要么跑得比num_workers=0还慢，要么直接报worker崩溃。原因很简单：多进程不是免费的，进程启动有开销，进程间数据传递有序列化消耗。我建议按这个节奏来调：先用num_workers=0跑通全流程，确认数据和模型逻辑没问题。图像数据规模中等（几千张），开num_workers=4通常够了。数据量大、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_memory=True会让DataLoader把数据放到锁页内存里，这个操作本身不影响CPU上跑的数据内容，但当你执行tensor.to(device, non_blocking=True)时，CPU到GPU的拷贝可以走异步路径，跟当前线程的其他计算重叠，从而稍微降低数据传输的阻塞时间。代价是锁页内存管理比普通内存更严格，占用过多会拖慢系统。drop_last=True则负责丢掉最后一组不够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(brightness=0.3, contrast=0.3, saturation=0.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_size=64, num_workers=4): train_ds = ImageFolderDataset( os.path.join(root_dir, 'train'), transform=train_transform ) val_ds = ImageFolderDataset( os.path.join(root_dir, 'val'), transform=val_transform ) train_loader = DataLoader( train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, drop_last=True, pin_memory=True, ) val_loader = DataLoader( val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True, ) 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_size=64) model = resnet18(num_classes=2).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) for epoch in range(10): model.train() for images, labels in train_loader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) 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_blocking=True) labels = labels.to(device, non_blocking=True) outputs = model(images) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) print(f"epoch {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', nrow=4, normalize=True, value_range=(-3, 3))打开aug_vis.png，如果里面图片内容都认不出来，多半是ColorJitter的幅度设大了或者RandomResizedCrop裁剪比例过狠，需要回调。4.4 磁盘类型对数据管道的影响：有人在这上面被坑惨了同样的代码、同样的num_workers，在机械硬盘和固态硬盘上的表现可能相差数倍。机械硬盘的随机读取能力很差，而图像数据集恰恰是大量小文件的随机读取，这在机械硬盘上是灾难性的。如果数据量不算特别大，一个实操有效的优化是：把图片打包成一个大文件存储，或者干脆在第一次遍历时缓存到内存。内存的读取速度比任何盘都快几个数量级，只要内存放得下，就把缓存空间利用起来。如果内存放不下，再退而求其次，把预处理后的特征事先算好存成本地.pt文件，训练时只读特征文件。这样做训练阶段直接绕开原图和transform，速度提升非常明显。代价是需要额外跑一遍离线预处理，且磁盘占用会变大。对“样本很大、重复训练很多轮”的场景，这笔投入很划算。5. 数据加载常见问题与排查实录5.1 高频问题速查表数据管道的坑看起来千奇百怪，实际上高度重复。把最常见的情况整理成一张表，遇到问题先对号入座。现象常见原因排查方向每个epoch拿到的第一个batch完全一样shuffle=False，或使用了sampler但没有手动打乱确认训练集shuffle=True，不要同时指定sampler训练时程序卡死无报错Windows下没加if __name__ == '__main__'保护把创建数据集和loader的代码全部放进main()里报错DataLoader worker exited unexpectedly子进程在getitem里抛异常，主进程看不到完整堆栈先设num_workers=0跑出完整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__返回的标签没转成int64，CrossEntropyLoss计算时类型不匹配。这种问题在打印labels的一瞬间就能发现。第二个必查点是增强后的可视化。save_image保存几张增强后的图片，肉眼看一遍，很多异常原因会比看代码更快暴露。比如RandomResizedCrop的scale参数如果设置不当，可能让图片被裁到只剩一个角；ColorJitter参数过大会让整个色调失真。这些“增强过度”的问题，不通过可视化很难从数值上感知。排查数据管道时，永远先怀疑数据，再怀疑模型。模型结构可以用预训练权重兜底，但错误的数据喂进去，再好的模型也只会越训越偏。5.3 排查思路：先把所有条件降到最保守档当遇到一个相当诡异、一时摸不着头脑的问题时，我的习惯是先把所有可能引入干扰的配置全部去掉。具体操作是：num_workers=0，只保留单进程加载。pin_memory=False，跳过锁页内存。shuffle=False，固定数据顺序。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, dtype=torch.long) padded_mask = torch.zeros(len(batch), max_len, dtype=torch.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, shuffle=True) train_loader = DataLoader( train_dataset, batch_size=64, sampler=sampler, num_workers=4, ) 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, desc=f'Epoch {epoch + 1}') for images, labels in pbar: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() pbar.set_postfix(loss=loss.item())如果GPU利用率偏低但进度条推得很慢，瓶颈多半在数据加载侧；如果GPU利用率很高但进度条依然不快，说明模型计算本身就重，数据管道已经喂得够快了。这个“看进度条判断瓶颈在哪端”的小技巧，我在实际项目里判断数据管道优化空间时，几乎是第一直觉手段。6.4 数据管道是地基，花时间打磨永远值得再分享一个我已经固化成习惯的流程。每接手一个新的数据集，我不会直接开足马力去调大num_workers或追求高batch_size，而是先写一个极简的inspect脚本，打印前几个batch的shape、dtype、数值范围，保存两张增强后的图片看一眼结果。等确认一切合理后，再进入正式训练循环。这个小习惯帮我避开了至少两位数的“训练跑完才发现数据有问题”的尴尬场景。如果你此刻正在一个新的PyTorch项目上起步，我建议你也把这个inspect环节固定下来。别急着把模型写多复杂，先把 Dataset 和 DataLoader 这一层打磨扎实。数据管道是深度学习中所有结论的地基，地基稳了，后面换模型、调损失函数、做实验对比，才不白折腾。