用PyTorch写自定义Dataset时，数据加载太慢怎么办？

遇到这种问题太正常了，10万张图全走JPEG解码加transforms，瓶颈根本不在GPU而在CPU那边。我建议你先别急着上num_workers，把worker设成2或者3试试，同时把persistent_workers=True加上，不然每个epoch反复创建进程开销巨大。内存炸多半是因为worker里每个都复制了完整的数据集引用，你可以试试把图片路径列表用共享内存或者直接把数据放到LMDB里，读取速度能快好几倍。至于transforms，随机裁剪和翻转确实会拖慢速度，但真正影响大的是每次decode和resize，你可以先把所有图片离线预处理成固定尺寸的numpy数组或uint8的tensor存成.pt文件，训练时直接load，这样能省掉大部分IO时间。还有个思路是给DataLoader加prefetch_factor，比如设成2，让每个worker提前多取几个batch，能缓解等待时间。最后关于tfrecord，PyTorch这边没有完全等价的东西，但你可以用WebDataset或者简单的HDF5，效果也不错。我自己的项目也是爬的图，后来改成离线转tensor加2个worker，epoch时间直接砍了七成，你可以试试。

我之前也踩过这个坑，10万张图全走jpg解码确实扛不住。建议你先试试把图片预处理成png或者直接存成npy/tensor格式，读的时候省掉大部分解码时间，内存不够的话就用lmdb或h5py分批存。另外num_workers炸内存不一定是worker太多，看看是不是pin_memory和persistent_workers没配合好，或者每个worker加载的batch太大。transforms里的随机操作其实影响不大，主要瓶颈在IO和解码，可以先用一个轻量级流水线测下瓶颈到底在哪。

内存炸了大概率不是num_workers本身的问题，而是每个worker都会复制一份完整的dataset对象，如果你在__getitem__里实时读图解码，10万张图全路径列表加上transforms的随机状态，每个worker额外占用几百MB很正常，4个worker叠加起来就爆了。我建议先检查一下是不是在dataset初始化时把整图都load进内存了，如果是的话改成懒加载，只在__getitem__里读单张图，这样内存占用会小很多。另外你说的把图片转成Tensor存起来，这个方向是对的，但别存成单个大文件，最好按类别分几个tar或者h5py文件，读取时用mmap模式，能省内存还能加速。transforms里的随机操作别去掉，但可以移到GPU上做，比如用torchvision.transforms.v2直接在CUDA张量上跑，或者用albumentations的GPU版本，这样CPU只负责解码，瓶颈会小很多。还有个土办法，如果你的图片尺寸统一，可以直接用torchdata的DataPipes，它支持流式读取和预取，比自定义Dataset更灵活。至于tfrecord，PyTorch这边有webdataset或者ffcv，效果类似，但配置成本有点高，你先试试把num_workers降到2，配合pin_memory=True和persistent_workers=True，看看能不能稳定跑起来，我怀疑你内存炸可能和shared memory设置有关，可以调一下/dev/shm的大小。最后如果还不行，就用torch.utils.data.IterableDataset配合多进程读取，但逻辑会复杂些，建议先把基础调优做完再考虑。

10万张图直接硬扛transforms确实遭不住，我当初也踩过这坑。建议你把预处理后的图像先存成.pt或者npy格式，训练时直接load tensor，能快好几倍。另外num_workers报错大概率是内存爆了，可以试试把persistent_workers=True加上，或者把batch_size调小点，别一次全塞进去。还有transforms里那些随机操作能少就少，尤其是随机裁剪这些，实在要留就放训练集里，验证集别加。

我之前也踩过这个坑，十万张图全走transforms确实顶不住。建议先把resize和归一化这种确定性操作提前做掉，存成预处理后的tensor或者npy，训练时只留随机增强，能快不少。另外num_workers报错大概率是workers之间复制了太多内存，试试把persistent_workers=True加上，或者把batch_size调小点，别一上来就4个worker。还有个思路是直接上WebDataset或者tar打包，流式读取，比单张读IO快很多，你可以看看这个方向。

我之前也踩过这个坑，10万张图全走transforms确实顶不住。建议把resize和归一化这类确定性操作提前做一次，存成预处理后的jpg或npy，训练时只读文件，随机增强留到Dataset里做就行。另外num_workers报错大概率是内存爆了，可以试试把persistent_workers和pin_memory打开，或者把worker数降到2，配合prefetch_factor调小一点。至于转Tensor，如果硬盘够快其实没必要，但你要是用机械硬盘，强烈建议先转成lmdb或h5py，IO能快好几倍。

说到这个我可太有感触了，之前也是自己爬数据做分类，10万张图差不多，刚开始跟你一模一样，一个epoch能跑到怀疑人生。你那个内存炸了的问题，我猜是num_workers开太高了，加上transforms里的随机操作会复制好几份数据，4个worker其实不算多，但如果你每张图都是原图读进来再resize，那内存肯定扛不住。我后来试了个笨办法，先把图片统一resize到256x256然后存成npy或者jpg，Dataset里直接读预处理好的文件，load速度能快好几倍，虽然占点硬盘空间但很值。transforms里那些随机翻转、裁剪之类的操作确实挺耗时的，我一般只在训练集用，验证集和测试集只做必要的resize和归一化。还有个小技巧，如果你显存够大，可以把预处理后的tensor直接缓存在内存里，用ramdisk或者简单的dict缓存，第一次读完之后后面就很快了。至于你说的tfrecord，PyTorch这边虽然没官方对应的格式，但可以用webdataset或者lmdb，效果也差不多，不过我觉得前期最有效的还是先把图片尺寸统一，减少每次load时的计算量。最后建议你检查一下是不是磁盘IO的瓶颈，如果图片是小文件的话，机械硬盘随机读就是慢，换SSD或者把数据打包成tar都能改善。

我之前也踩过这个坑，十万张图其实不算多，瓶颈多半在IO和解码上。你试试把图片预处理成pin_memory=True，配合num_workers=4但把persistent_workers也打开，内存炸可能是worker数超过物理核心了。另外强烈建议先离线把所有图片resize成统一尺寸再存成.pt或LMDB，训练时只做ToTensor和归一化，随机操作放训练循环里单独做，速度能快好几倍。至于tfrecord，PyTorch这边可以用WebDataset或者简单的tar打包，效果类似但更轻量。

10万张图一个epoch两小时确实有点离谱了，我怀疑瓶颈不一定全在transforms上。你试试把num_workers设成2或者6，有时候4反而会卡在内存带宽上，尤其是如果你机器物理内存不大还开了pin_memory的话，很容易直接OOM。另外你提到把图片转成Tensor存起来，这个方向是对的，但别存成单个大文件，我建议你按类别分几个子集，用lmdb或者h5py打包，读取时用内存映射，速度能快好几倍。transforms里那些随机操作确实挺吃CPU，但如果你只是resize和归一化，其实开销还好，真正慢的往往是decode，所以可以考虑先把图片缩到较小的尺寸再存，比如256x256，训练时再随机crop。还有个偏门点的方法，就是干脆不用torchvision的Dataset，自己写个多进程预取队列，把图像读取和预处理放到单独的进程池里，主进程只拿结果，这样能绕开DataLoader的一些锁开销。最后，如果你不想折腾存储格式，也可以试试NVIDIA的DALI，它对JPEG解码有专门优化，配合GPU做预处理，10万张图应该能压到几分钟一个epoch，不过配置起来有点学习成本。

10万张图确实得上点手段了，我之前也踩过这坑。你transforms里如果每次都在线做resize和归一化，CPU肯定扛不住，尤其jpeg解码本身就慢，这俩活儿全挤在主进程里了。我建议你先别急着上num_workers，把worker数从2开始慢慢调，同时把persistent_workers=True加上，不然每个epoch都要重新创建进程，内存反而容易炸。另外你说的存Tensor是个好思路，可以先用脚本把所有图预处理成.pt或者.npy文件，训练时直接load，省掉解码和resize的耗时，代价是磁盘占用会大不少，但10万张图应该能接受。至于transforms里的随机操作，像随机裁剪、翻转这些尽量别砍，它们对精度有帮助，但你可以把它们移到GPU上做，比如用torchvision.transforms.v2，或者干脆在Dataset里只做轻量级的归一化，重活留给GPU的collate函数。还有个小技巧，如果你用Linux，可以考虑把数据放到/tmp或者用tmpfs做内存盘，IO速度能快一个量级。不过说实话，最省事的方案还是先试试DataLoader里加个prefetch_factor=2，再把num_workers调到你的CPU核心数一半，很多时候瓶颈不在解码，而是进程间通信调度太频繁了。

我最近也踩过这个坑，10万图跑俩小时确实有点离谱了。你那个内存炸了的报错，大概率是num_workers开的进程太多，每个worker都会复制一份dataset的引用，加上transforms里的随机操作会额外占用内存，可以先降到2试试，或者把prefetch_factor设小一点。不过更推荐的做法是先把图片预处理成张量存下来，比如全部resize到统一尺寸后存成.pt文件或者npy格式，训练时直接load张量，省掉每次读图解码的IO开销，我试过能快三四倍。transforms里的随机增强其实影响不大，瓶颈主要在磁盘读取和JPEG解码上，如果你用普通硬盘的话，建议先试试把图片转成lmdb或者h5py这种二进制格式，顺序读取会快很多。至于tfrecord，PyTorch这边没有官方等价物，但可以用webdataset这个库，它跟tfrecord思路差不多，把样本打包成tar文件流式读取，配合num_workers效果不错。另外你如果显存够的话，可以试试把整个预处理后的数据集直接load到内存里，10万张resize到224的话大概20GB左右，一般服务器能扛住，这样训练速度直接起飞。最后检查下你是不是在train的时候又把transform跑了一遍，有时候验证集和训练集共用同一个dataset类会重复预处理，写个flag区分一下也能省不少时间。

我之前也踩过这个坑，10万图其实不算特别多，但transforms里如果每次都在线做resize和归一化，CPU直接成瓶颈。你可以试试先把图片预处理成uint8的numpy或者直接存成.pt的tensor，训练时只做ToTensor和归一化，速度能快好几倍。另外num_workers报错大概率是每个worker都会复制一份数据集引用，内存翻倍，可以试试把prefetch_factor调小一点，或者用persistent_workers=True看看。至于tfrecord，PyTorch这边可以用WebDataset或者自己写个LMDB，但前期转换也挺费时间的，不如先优化预处理缓存来得直接。

10万张图还带transforms，单epoch两小时确实有点离谱了，我猜你瓶颈大概率不在CPU预处理，而是磁盘IO和GPU喂数据之间的衔接出了问题。num_workers设4就爆内存，可能是你transforms里用了RandomResizedCrop这类会临时复制多份数组的操作，每个worker都开一份完整图像副本，内存直接翻倍，建议先试试把worker数降到2，同时把persistent_workers=True加上，能省掉反复创建进程的开销。另外你说把图片转成Tensor存起来，这思路可行但别直接存成pkl或npy，因为文件体积会膨胀到几个G，加载时反而更慢，我建议先用jpg或png把预处理后的图落盘成固定尺寸，然后用lmdb或h5py做键值存储，读取时按索引随机访问，比遍历小文件快很多。至于transforms里少用随机操作，这个对提速帮助有限，因为随机操作本身计算量不大，真正吃时间的还是resize和归一化里的像素级循环，你可以考虑把归一化挪到GPU上做，比如在collate_fn里只做resize，然后进模型前在cuda上用F.normalize。PyTorch没有完全对应tfrecord的东西，但有个webdataset库，能把10万张图打包成tar文件，训练时流式读取，配合num_workers能接近线性加速，你可以试试看。最后建议你profile一下，用torch.utils.bottleneck看看到底是数据加载耗时占比高还是模型前向耗时高，别盲目优化，我之前就是没查清楚，折腾半天才发现是数据增强里用了太多PIL的resize导致CPU占用满了。

10万张图全走实时transforms确实扛不住，尤其resize和归一化这种CPU密集操作，4个worker内存炸多半是每个worker都复制了一份完整的数据索引或者缓存，建议先检查下是不是把整个图片列表传进去了。我之前遇到类似情况，是把所有图片预处理成256x256的RGB uint8数组直接存成npy或者lmdb，训练时只做ToTensor和归一化，速度能快三倍以上。不过要注意如果transforms里有随机裁剪或翻转，得在读取时动态做，但尽量用GPU上的随机操作，别在Dataset里做。另外你也可以试试把图片先缩到一个小尺寸缓存到内存里，比如用pickle存成list，这样每个worker加载时内存占用会小很多。至于tfrecord，PyTorch这边有WebDataset或者简单的HDF5方案，但我觉得对你这个规模，lmdb或者npy可能最省事。还有个小技巧，把num_workers设成0跑一次看下瓶颈是不是真在IO上，如果0也很慢那可能是transforms本身太重了，可以考虑用albumentations这类加速库。最后提醒下，如果内存还是不够，可以试试把数据分几个文件夹，每个epoch只加载一部分，但这样会打乱全局shuffle，得看你对精度要求高不高。

我之前也踩过这个坑，十万张图全放内存确实容易爆，但你可以试试把图片预处理后存成lmdb或h5py格式，读起来比直接读文件快很多，而且内存占用可控。transforms里的随机操作其实影响不大，瓶颈主要在IO和decode上，建议把resize和归一化提前到保存时做，训练时只做随机裁剪或翻转。另外num_workers不是越高越好，我一般先用2测一下内存峰值，再逐步加，同时配合persistent_workers=True能减少重复初始化开销。如果你不想折腾，也可以看看torchdata这个库，它支持类似tfrecord的流式读取，但生态还没那么成熟。

看到你这个描述我太有同感了，之前自己爬数据做训练也是卡在这一步。内存炸了多半不是worker数量的问题，而是每个worker都会复制一份完整的Dataset对象，如果你在__getitem__里实时读图再transform，那4个进程同时做JPEG解码和resize，内存当然扛不住。我的建议是别急着把所有图片转成Tensor存硬盘，那个文件会巨大而且读取也不一定快，更好的做法是做一个预处理缓存，比如第一次跑的时候把resize和归一化后的结果存成npy文件，之后每次直接load npy，这样能省掉大量重复计算。另外你提到transforms里的随机操作，其实像RandomCrop这种如果每次都做确实很耗时，可以先离线把图裁好存下来，在线只做ToTensor和Normalize这种轻量操作。还有个坑是图片格式，如果你爬下来的图是PNG，转成JPEG后解码速度会快不少，文件也小。最后如果还不能解决，可以试试把数据打包成HDF5或者LMDB，PyTorch虽然没有官方tfrecord那么方便，但用这两个格式配合DataLoader的persistent_workers和prefetch_factor调参，速度能明显上来。

跟你遇到一模一样的问题，10万张图raw读确实扛不住。我后来是把所有图片预处理成png或jpg的tensor存成单个.pt文件，用torch.load一次性加载进内存，训练时直接切片，速度提升非常明显，内存大概占10G左右，你可以试试。另外num_workers报错不一定是内存本身，很可能是你用了pin_memory=True加上worker数量太多，系统共享内存不够，改成2个worker然后pin_memory=False看看。transforms里的随机操作不影响加载速度，但如果你每次epoch都做随机resize，确实会让CPU压力大，建议把resize这类重操作放到预处理阶段一次性做掉，训练时只留归一化和随机翻转这种轻量的。至于tfrecord，PyTorch生态里对应的是webdataset或者lmdb，但说实话对于10万张图，先转成.pt或者h5py就够了，没必要上那么重的方案。还有个坑是如果你的图片尺寸不统一，别用default_loader，自己写个读取函数直接cv2.imread，配合albumentations做预处理，比torchvision的transforms快不少，因为后者在PIL和Tensor之间反复转换。最后建议你做个缓存，把每个样本的预处理结果存成npy，第二次跑直接读npy，这样最省事。

我之前也踩过这个坑，10万张图全走jpg解码确实扛不住。建议你先试试把图片预处理成png或者直接存成npy/tensor格式，读取速度能快好几倍，内存不够的话可以分块存。num_workers报错多半是shared memory不够，可以试试把worker数降到2，或者用persistent_workers=True，再配合pin_memory=True，一般能缓解。另外transforms里那些随机操作别放CPU端，能转到GPU上做就转到GPU上，尤其是resize这种大头。

建议先用脚本把图片预处理后存成npy或lmdb，读取时只做张量化，能快好几倍。另外worker别乱加，内存不够就调低persistent_workers试试。

我之前也踩过这个坑，10万图全走transforms确实扛不住。你试试把resize和归一化直接写进一个离线预处理脚本，一次性把图片转成npy或者jpg缓存下来，训练时Dataset里只做ToTensor，能快好几倍。num_workers报错的话，把persistent_workers=True加上，然后worker数量从2开始慢慢调，别一上来就4，内存不够就换prefetch_factor=2。另外随机操作少用点确实有影响，但主要还是IO瓶颈，先解决缓存再说。