最近在做一个图像分类项目,数据集是自己爬的,大概有10万张图片。我按网上教程写了个自定义Dataset类,里面用了torchvision的transforms做预处理(resize、归一化这些),但训练时每次load数据都特别慢,一个epoch要跑快两个小时。我看别人说用DataLoader的num_workers能加速,但我设了4个worker之后反而报错了,好像是内存炸了。想问下各位大佬,这种自定义数据集的情况,有没有什么常规的优化技巧?比如是不是应该先把图片转成Tensor存起来,还是说transforms里少用点随机操作?或者PyTorch有没有像TF那样可以直接从tfrecord读数据的方式?先谢过大家了。
楼主
2026-07-19
用PyTorch写自定义Dataset时,数据加载太慢怎么办?
请 登录 后发表回复
全部回复
共 162 条
2楼
2026-09-26
10万张图每轮两小时,瓶颈大概率在CPU解码和resize上,而不是IO。你设num_workers=4内存炸,多半是每个worker都复制了一份Dataset对象,里面要是有大缓存就直接翻倍了。可以试试先把图片统一resize到短边256左右存成jpeg或lmdb,训练时再随机crop,能省不少解码开销。另外transforms里的随机操作本身不慢,慢的是PIL解码,换成opencv或turbojpeg会快很多。
3楼
2026-09-26
10万张图每epoch两小时,瓶颈大概率在IO和CPU预处理上。我建议先用LMDB或webdataset把图片和标签打包成二进制,读的时候直接内存映射,比小文件随机读快非常多。num_workers开到4内存炸,通常是每个worker都复制了一份数据或者transforms里有大缓存,可以试试persistent_workers加prefetch_factor,再把resize这类固定操作提前离线做好。另外torchdata的DataPipe现在支持类似tfrecord的流式读取,值得看一下。