PyTorch数据管线实战：Dataset与DataLoader从入门到调优-代码聚汇网

PyTorch数据管线实战：Dataset与DataLoader从入门到调优

PyTorchDatasetDataLoader

|难度：普通|阅读约 49 分钟

如果你用PyTorch做过图像分类，大概率遇到过这些场面：数据集几万张图，一股脑塞进内存直接爆掉；或者训练集忘了打乱，模型训了半天loss不降反升；再或者num_workers调了个大数值，程序刚启动就报错崩溃。这些问题的根子，基本都出在Dataset和DataLoader这两个最底层、最高频的组件上。

Dataset和DataLoader是PyTorch数据管线的核心。一个负责回答“我有哪些数据、如何取出一条样本”，另一个负责回答“训练时怎么才能高效、稳定、有序地把数据一批批喂给模型”。这篇文章不打算讲花哨的底层源码，而是从实际使用经验出发，把这两个组件的分工、实现方式、参数取舍、常见坑和性能优化思路完整梳理一遍。适合刚上手PyTorch的初学者，也适合那些明明网络模型没问题、却被数据加载搞到怀疑人生的老哥。

1. 先搞清楚：Dataset和DataLoader这对搭档到底在解决什么问题

1.1 没有数据管线的年代，数据加载有多痛

很多初学者刚接触PyTorch时，第一反应是“我直接把图片读进list里不就行了？”。如果你的数据集只有几百张小图，确实可以这样干。但一旦进入真实项目，画面就完全变了：数据量动辄几万、几十万，甚至上百万，每张图片分辨率还不小，全部读进内存，16G内存可能都装不下。就算内存勉强够，手动写训练循环时还得自己处理切片、打乱、批次化、归一化，代码又长又容易出错。

举个最典型的场景：你没有用DataLoader，手动按batch训练。

python复制# 原始方式：手动切batch，还要自己处理顺序
indices = list(range(len(data)))
random.shuffle(indices)
for i in range(0, len(indices), batch_size):
    batch_indices = indices[i:i+batch_size]
    batch_data = torch.stack([data[j] for j in batch_indices])
    batch_label = torch.tensor([label[j] for j in batch_indices])
    # ...然后才开始训练

这套流程有几个硬伤。第一，所有数据必须预先加载成Tensor，大数据集内存直接爆。第二，shuffle、batch切分、Tensor转换全要自己写，代码没法复用。第三，完全没有多进程并行能力，数据加载和模型训练串行执行，GPU经常在那等数据，利用率上不去。

Dataset和DataLoader就是为了解决这套麻烦而存在的。它们的核心思想是把“数据怎么管理”和“数据怎么取用”两件事拆开，各管一摊，同时把并行加载、自动打乱、批量拼接这些高频操作内置好，让开发者把精力放在模型本身。

1.2 分工：一个是数据源，一个是搬运管线

我习惯用一个生活化的类比来解释两者的关系。Dataset就像超市的仓库货架，每一格都摆着一个样本，你告诉它“给我第3个位置的货物”，它就把那件货物取出来给你。它不关心你要不要批量买、按什么顺序买。DataLoader则是超市里的收银台和购物车，负责把仓库里的货物按照你的要求（每批装几件、要不要打乱顺序、用几个员工同时搬运）整理成一批批的购物袋，送到你面前。

落到代码层面，Dataset的核心是三个方法：

__init__：初始化路径列表、transform等配置
__len__：返回数据总量
__getitem__：按索引返回一条样本（通常是(输入, 标签)）

DataLoader则是在Dataset外面套了一层调度逻辑。它内部会调用__getitem__取出样本，然后按照batch_size合并成一个batch，再通过collate_fn把多条样本拼成带batch维度的Tensor。

掌握这个分工后，你就明白了一条核心原则：不要在一开始的Dataset里就把所有数据load到内存，也不要让DataLoader去管数据从哪里来。数据源和数据管线的职责一旦混在一起，后续的扩展、换数据集、调优都会变得束手束脚。

需要模型API调用？免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。

2. 手写一个Dataset其实没那么玄乎：三种主流实现全给出

2.1 自定义Dataset的标准写法：重写三个方法就够了

实战里最常用的就是自定义Dataset。以图片分类为例，我通常会把样本路径存成一个列表，在__getitem__里读取图片、做transform、返回Tensor。

python复制import torch
from torch.utils.data import Dataset
from PIL import Image
import os

class ImageFolderDataset(Dataset):
    def __init__(self, file_list, label_list, transform=None):
        self.file_list = file_list
        self.label_list = label_list
        self.transform = transform

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        img_path = self.file_list[idx]
        label = self.label_list[idx]
        image = Image.open(img_path).convert('RGB')
        if self.transform:
            image = self.transform(image)
        return image, label

这个写法有几个关键细节值得注意。

第一，为什么要把transform放在__getitem__里而不是__init__里？因为一张图在__getitem__被调用时才真正读出来、做增强，同一张图每次被取到时可能做不同的随机变换，这符合训练集数据增强的预期。如果在__init__里一次性把所有图都处理好，几百G数据没等训练就先把内存吃干净了。

第二，路径列表最好在__init__里就全部准备好。踩过坑的人都知道，如果在__getitem__里才去遍历目录找文件，每次取样本都会做一次文件扫描，速度慢到怀疑人生。正确做法是在初始化阶段扫描好目录，把每条样本的地址和标签存成list，__getitem__只负责按索引取。

第三，__len__必须返回准确的样本总数。DataLoader的很多逻辑（比如进度条、shuffle、sampler）都依赖这个数字，写错了会出现数据取不到头或者索引越界的问题。

我在实际项目里经常会在Dataset里额外加一个get_filename之类的辅助方法，方便出问题时定位是哪张图导致训练异常，这个习惯排查bug时非常有用。

2.2 不想手写就用内置方案：ImageFolder和torchvision.datasets

如果数据集的目录结构恰好是根目录/类别名/图片这种形式，可以不用手写Dataset，直接用torchvision.datasets.ImageFolder。

python复制from torchvision.datasets import ImageFolder
from torchvision import transforms

train_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

train_dataset = ImageFolder('./data/train', transform=train_transforms)

ImageFolder会自动把子目录名作为类别名，并按字母序映射成数字标签。dataset.classes可以查看类别名列表，dataset.class_to_idx可以查看映射关系。这个类的优点是省事，但代价是灵活度有限。比如你想根据一个CSV表格记录来指定每张图的标签，或者要做多标签分类，ImageFolder就不好使了，这时候还是乖乖回到自定义Dataset。

torchvision.datasets下面还有很多内置数据集，比如CIFAR-10、MNIST、ImageNet接口。以CIFAR-10为例：

python复制from torchvision.datasets import CIFAR10

train_dataset = CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform=train_transforms
)

内置数据集的价值不只是省事，关键是可以和别人的实验结果对齐——同样用CIFAR-10、同样的预处理，大家跑出来的指标才有可比性。

2.3 TensorDataset：纯Tensor数据的快速方案

如果你的数据已经是Tensor形态（比如从numpy读进来、或者特征已经抽取好了），那就没必要走文件读取的路。直接用TensorDataset，几行代码就能建好Dataset。

python复制from torch.utils.data import TensorDataset
import torch

features = torch.randn(10000, 128)   # 10000条样本，每条128维
labels = torch.randint(0, 10, (10000,))  # 10分类
dataset = TensorDataset(features, labels)

TensorDataset本质是按索引同时取各个Tensor的对应行，然后组合成(feature, label)返回。它不能做transform，也无法动态处理增广，所以适合文本特征、表格数据这类已经预处理好的场景。你要是手头有numpy数组，可以先torch.from_numpy(...)转成Tensor再喂给TensorDataset。

这个方案的另一个好处是调试方便。我想快速验证一个模型能不能跑通，就用 TensorDataset 造一份假数据，先把训练流程跑通，再替换成真正的数据集，效率非常高。

3. DataLoader参数详解：8个参数管住你的数据管线

3.1 batch_size怎么定：不只看显存，还要看模型类型

很多新手问：batch_size到底设多少？我给的答案是——先看显存能装下多少，再看你的模型对batch size的敏感度。

显存方面可以用一个简单的估算方法：先随便设一个batch_size（比如16），跑一次前向传播，观察显存占用，然后按比例推算。如果你的模型在batch=16时占用显存6G，那你用8G显存的卡，batch_size最多也就20出头，再往上就会OOM。

模型类型也要考虑。BatchNorm层的效果和batch里的样本统计量强相关，batch太小（比如小于8）统计量不稳定，训练容易震荡。目标检测里经常看到batch=2、batch=4这种配置，那是因为显存实在有限，那就得配合gradient accumulation来模拟大batch效果。文本分类里的动态padding也需要考虑batch内最大长度的影响。

推荐的做法是：图像分类这类任务，先从32开始，结合显存和训练曲线调整；如果显存充足，可以尝试64或128，但要注意学习率通常也需要跟着调。变长输入的任务（NLP、语音），batch先小一点，给padding留足计算空间。

3.2 shuffle到底管的什么：训练集要打乱，验证集别打乱

shuffle这个参数非常基础，但很多人理解不深。简单说，shuffle=True表示每个epoch开始前把数据顺序打乱，这样每个batch里的样本组成都不同，避免模型学到数据顺序带来的假规律。

如果训练数据本身有排序性（比如前5000张都是猫，后5000张都是狗），不打乱就会导致一个batch内全是同一类，模型训练的梯度来回震荡，loss曲线跟锯齿一样难看。很多人遇到loss不下降、acc纹丝不动，排查半天，最后发现是忘了开shuffle。

验证集和数据测试集则要设置shuffle=False。因为验证时我们想要的是稳定的、可复现的评估结果，不需要打乱顺序。同时验证集通常配合torch.no_grad()，不打乱还能方便你把预测结果和原始样本对齐，排查错分样本时非常有用。

python复制train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

3.3 num_workers怎么选：不是越大越快

num_workers控制DataLoader用几个子进程来加载数据。默认值是0，意思是数据加载就在主进程里做，好处是不会有多进程的额外开销，坏处是GPU在训练时必须等数据读好，很多时间都浪费在等待上。

设置num_workers>0后，DataLoader的子进程会提前把后续batch的数据准备好（prefetch），GPU训练当前batch时，子进程已经在后台读下一个batch了，训练和加载实现了流水线并行。

但这不意味着num_workers越大越好。每个worker会额外占用CPU和内存，worker数量过多时，CPU调度开销反超收益，系统甚至可能因为争抢内存而崩溃。我常用的策略是：先看机器的物理CPU核心数，num_workers从4或8起步，观察训练时的CPU占用和GPU利用率。如果GPU利用率（nvidia-smi）已经稳定在90%以上，说明加载速度不是瓶颈，继续加worker没有意义。如果GPU利用率偏低（70%以下）且CPU还有很多余量，就逐步调大num_workers，通常收益明显。

一个我在实践中总结的保守经验：worker数不要超过CPU物理核心数，更稳妥的是用核心数的一半。比如8核CPU先设4，跑起来看情况再调。

3.4 collate_fn：批处理拼接逻辑的“自定义开关”

collate_fn是DataLoader里最容易被忽略、但关键时候最管用的参数。它的作用是把Dataset返回的一批样本拼成一个batch。默认的collate逻辑会把每个样本的Tensor叠到第0维，形成(batch_size, ...)的新Tensor。但遇到变长文本、变长语音这类样本时，默认拼接会直接报错。

比如NLP里每条句子长度不同，直接stack会维度对不上，这时候就要自己写collate_fn，在batch内做动态padding，统一成当前batch的最大长度。

python复制def collate_batch(batch):
    inputs = [item[0] for item in batch]
    labels = torch.tensor([item[1] for item in batch])
    lengths = torch.tensor([len(x) for x in inputs])
    max_len = lengths.max().item()
    padded = torch.zeros(len(inputs), max_len, dtype=torch.long)
    for i, x in enumerate(inputs):
        padded[i, :len(x)] = x
    return padded, labels, lengths

有了lengths，后续做RNN或Transformer时可以直接配合torch.nn.utils.rnn.pack_padded_sequence或者attention mask。collate_fn还有一个隐藏用法：把读取到的PIL Image统一转换成Tensor、做归一化。有些项目把transform放在collate里而不是Dataset里，也能跑，但我的习惯是transform尽量放Dataset，collate只做结构上的批处理，这样职责更清晰。

3.5 pin_memory、drop_last和prefetch_factor的取舍

pin_memory=True的意思是，当数据在CPU内存中时，提前把它锁页，这样GPU拷贝数据时走更快的传输通道。实测下来，在GPU训练场景下，pin_memory=True通常能带来5%-15%的训练速度提升，尤其数据量大的时候更明显。代价是多占一点内存，但相比收益，通常值得。我的习惯是只要用GPU训练就设True，用CPU训练时保持默认False。

drop_last表示当样本总数不能被batch_size整除时，最后一个不完整的batch要不要丢掉。你会问：丢掉不是浪费数据吗？但在某些场景下必须丢——比如使用BatchNorm时，一个只有3条样本的batch会导致计算出的mean和variance极不稳定，影响模型效果。所以如果总样本数是1000，batch_size=32，1000除以32余8，drop_last=True会让每轮实际训练数据变成992条，丢掉最后8条。对于验证集，我通常不设drop_last，因为验证是逐个batch累计算指标，不完整batch没关系。

prefetch_factor是个小众参数，默认值是2，意思是每个worker预加载2个batch。在数据加载确实是瓶颈的前提下，可以把prefetch_factor调到4或8，减少worker等待时间。这个参数配合persistent_workers一起用，效果更明显。

python复制train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    drop_last=True,
    prefetch_factor=4,
    persistent_workers=True
)

persistent_workers=True表示每个epoch结束后不销毁worker进程，避免反复创建进程的开销。对于多epoch训练，这个参数的收益肉眼可见。但如果单次epoch数据量很小，worker还没捂热就训完了，持久化worker的意义就有限，反而增加内存占用。

4. 猫狗分类实战：把Dataset和DataLoader串成一套完整流程

4.1 先规划好数据和文件目录

说一千道一万，不如跑一遍完整流程。我们设计一个猫狗分类任务：数据目录结构如下样例。

code复制data/
  train/
    cat/
      cat_001.jpg
      cat_002.jpg
    dog/
      dog_001.jpg
      dog_002.jpg
  val/
    cat/
      cat_001.jpg
    dog/
      dog_001.jpg

这里直接用ImageFolder就能建出训练和验证Dataset，后面我再展开从CSV构造自定义Dataset的版本。先明确一点：ImageFolder按子目录名映射标签，所以目录建得规整，后面就省很多事。

4.2 构建Dataset和DataLoader的完整代码

这部分我给出一个图像分类任务里用得最多的标准模板，包含数据增强、规范化、加载器创建三个环节。

python复制import torch
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms

train_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

val_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

train_dataset = ImageFolder('./data/train', transform=train_transforms)
val_dataset = ImageFolder('./data/val', transform=val_transforms)

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,
    pin_memory=True,
    drop_last=True
)

val_loader = DataLoader(
    val_dataset,
    batch_size=32,
    shuffle=False,
    num_workers=4,
    pin_memory=True
)

这里有个细节：训练集做了随机水平翻转、随机旋转和颜色扰动，但验证集只做Resize和ToTensor。数据增强的目的是让模型见到更多样性的输入，提升泛化能力，但验证集需要稳定的评估口径，不能引入随机性，否则同一张图每次验证结果都不一样，指标就失去参考意义了。

4.3 在训练循环里正确使用loader

有了loader，训练循环就清爽很多。只需要一个for循环。

python复制device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(10):
    model.train()
    total_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')

    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Validation Acc: {100 * correct / total:.2f}%')

训练加载器因为shuffle=True，每次epoch取到的batch划分都不一样，模型不会记住训练样本的出现顺序。验证加载器不shuffle，因此每个epoch计算出来的acc是固定可复现的，调参时比较不同模型或不同超参数的效果才有意义。

这类训练循环还有个常见优化点：把.to(device)尽可能早做，而不是在模型forward里做。数据一取出来就转设备，后续计算都发生在GPU上，不会反复触发CPU到GPU的拷贝。如果你用TensorDataset加载假数据在CPU上跑通了，换成真实数据后只需要把路径和transform替换掉，剩下都不用动，这也是分层设计带来的好处。

5. 高频翻车现场：5个我踩过的DataLoader坑与排查思路

5.1 Windows下num_workers报错：不是代码问题，是进程模型问题

在Windows上使用num_workers>0时，经常遇到一个报错：RuntimeError: An attempt has been made to start a new process before the current process has finished its bootstrapping phase。

这个报错的本质是Windows和Linux的进程创建方式不同。Linux用fork，子进程直接复制父进程内存，所以DataLoader的子进程可以轻松拿到Dataset对象。Windows用spawn，子进程需要重新导入主模块来重建环境，如果你没有把训练代码放在if __name__ == '__main__':保护块里，子进程在导入主模块时会再次执行训练逻辑，无限递归下去，最终崩溃。

解决方案是在入口处加上保护：

python复制if __name__ == '__main__':
    main()

所有涉及DataLoader多进程的代码都要放进main函数。这句话是Windows下PyTorch训练的必背口诀。如果你遇到了，检查两件事：一是代码是否被主模块保护；二是是否在用Jupyter Notebook，如果是，num_workers只能设0或1，因为Notebook的交互式环境无法正确处理多进程spawn。

5.2 内存溢出：Dataset设计不当是主因

训练时内存一点一点涨，最后直接Out of Memory，这个问题我在初学阶段踩过多次。最常见的错误是在__init__里把所有图片都读成PIL图像对象存在list里。一张图几十到几百KB，一万张图就可能吃掉好几个G内存，而且这些对象不会被GC及时回收，内存只会越占越多。

正确做法是__init__里只保存文件路径，在__getitem__时才真正读图片。这样内存中同一时刻只有当前batch的图片，训练几万张图的内存压力和训练几百张图基本一个量级。

如果你用的是pin_memory=True，内存占用会额外增加一些，因为它需要锁页内存。如果内存比较紧张，可以先关掉pin_memory观察一下。还有一种情况是DataLoader的prefetch机制，num_workers和prefetch_factor越大，预取的数据越多，内存占用越高。内存吃紧时优先调低num_workers，再考虑prefetch_factor。

5.3 数据类别不均衡：Sampler比手动过采样好用

分类任务里经常遇到正负样本比例悬殊，比如猫狗数据里猫几百张、狗几万张。如果不处理，模型会严重偏向多数类。很多人第一反应是做数据复制，但更优雅的方案是用PyTorch的WeightedRandomSampler。

核心思路是给每个样本一个采样权重，少数类的样本被抽到的概率更高。权重通常设为1/样本数，再归一化。

python复制from torch.utils.data import WeightedRandomSampler

labels = train_dataset.targets  # ImageFolder里有targets属性
class_counts = torch.bincount(torch.tensor(labels))
weights = 1.0 / class_counts[labels]
sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    sampler=sampler,
    num_workers=4
)

这里要特别注意：一旦传入sampler，就不能再设置shuffle=True，因为sampler内部自己实现了抽样顺序，两者冲突时会直接报错。replacement=True表示允许同一张样本在一个epoch内被重复抽到，这能有效缓解少样本类别的数据不足问题。需要注意，带sampler的loader会按权重抽取指定数量的样本，所以len不再等于数据集长度除以batch，而是等于num_samples除以batch_size。

5.4 说不清IterableDataset和MapDataset的区别

PyTorch里有两种Dataset类型：MapDataset（最常用的自定义Dataset）和IterableDataset。前者是“按索引查表”，后者是“流式读取”。很多人把它们混着用，导致报错时很困惑。

MapDataset必须实现__getitem__和__len__，适合所有样本能随机访问的场景。IterableDataset只需要实现__iter__，继承自torch.utils.data.IterableDataset，适合无法预知长度、需要流式读取的场景，比如从数据库查询、实时读取传感器数据、从大型分布式文件系统按顺序读取等。

python复制from torch.utils.data import IterableDataset

class MyIterableDataset(IterableDataset):
    def __iter__(self):
        for i in range(1000):
            yield i, i * 2

IterableDataset和DataLoader的shuffle参数天生不兼容，因为shuffle需要先知道全量索引，而IterableDataset可能连长度都不知道。很多人把shuffle=True传给IterableDataset然后报错，实际上应该在Dataset内部自己维护一个buffer来做打乱，或者接受这种数据天然无法全局打乱的现实。

5.5 验证集loss曲线诡异：shuffle和drop_last的连锁反应

最后说一个容易忽略的细节。如果你在验证集上开了drop_last=True，且总样本数恰好不能被batch_size整除，那最后一批样本被丢掉，验证集指标每轮都在基于不同数量的样本计算，小数点后几位会出现细微跳动。这种波动虽然很小，但在调试超参时足以干扰你的判断。

我的习惯是：训练集用drop_last=True（避免BatchNorm受不完整batch影响），验证集和测试集用drop_last=False。如果为了严格对齐样本数而必须在验证集也drop_last，确保所有对比实验用完全一样的设置，否则结论不可靠。

6. 进阶调优：把数据加载速度再压榨一档的几种思路

6.1 预处理前置：把Image Decode从训练循环里搬走

很多人训练慢，瓶颈不是GPU，而是CPU在频繁做图片解码和Resize。如果你的图片是高清大图，每次__getitem__都要先解码、再Resize到224，耗时非常可观。一个重要思路是把数据预处理前置，在正式训练前把图片统一解码并缩放成小尺寸，存成压缩格式或npy数组。

比如可以先跑一个预处理脚本，把所有图片Resize到256x256并保存为jpg或.npy。训练时Dataset直接读取已经缩放好的图片，省去每次解码大图的开销。这在数据量几个G的规模下非常有效，训练速度可能提升20%-40%。

不过要小心一个坑：前置Resize会损失一部分原始信息，如果你的任务需要对原始图像做精细分析（比如目标检测的小目标），过早缩小图片可能影响最终精度。一般建议先Resize到比网络输入稍大的尺寸（比如256，网络输入224），给后续数据增强留一些裁剪空间。

6.2 用persistent_workers和更大的prefetch_factor榨干CPU

前面提过persistent_workers，这里展开说。默认情况下，每个epoch结束，DataLoader会销毁worker进程，下个epoch重新创建。对于小数据集，这个过程可能无所谓；对于大数据集，反复创建进程的时间累加起来非常可观。

python复制train_loader = DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,
    pin_memory=True,
    prefetch_factor=8,
    persistent_workers=True
)

几个参数配合上之后，worker会在内存里常驻，预取数据管线不会在epoch交界处断流。如果你的系统内存充足（比如32G以上），prefetch_factor调到8没什么问题。内存只有16G时先谨慎，prefetch_factor和num_workers一起增大会迅速推高内存占用。

6.3 混合精度与数据管线的整体协同

最后说一个经常被忽略的协同关系：当你用混合精度训练（torch.cuda.amp）加速模型计算时，GPU的算力余量变大，数据加载更容易成为瓶颈。我遇到的情况是，模型计算时间缩短了，但总训练时间没有按比例下降，用nvidia-smi一看，GPU利用率只有60%多，说明数据喂不饱GPU。

这时候优先调num_workers和prefetch_factor，把数据加载速度提上来。如果还不行，考虑在Dataset里把图片直接读取成RGB字节数组，减少PIL Image的额外开销；或者把图片格式从PNG换成JPEG，解码速度更快。我自己做过一个对比，同一任务下PNG转JPEG后，数据加载耗时下降了约25%，代价是略微的编码质量损失，但训练任务里完全可以接受。

所谓整体协同，就是不要只盯一个环节。先看GPU利用率，再定位瓶颈在模型计算还是数据加载，最后针对性地调参。数据管线和模型训练是一个流水线，最慢的环节决定总速度。

我在实际项目里的体会是：Dataset和DataLoader的门槛不高，但把所有参数吃透、把各种边界情况处理好，需要不少实战积累。数据加载这块的问题往往不是一次报错就能发现的，它们会潜伏在训练过程中，表现为GPU利用率低、loss波动大、验证指标不稳定。建议拿到新项目时，先花半小时把Dataset和DataLoader单独跑通，打印一下每个batch的shape、类型、数值范围，确认无误后再开始训练模型，能省下后面大量排查问题的时间。