pytorch dataloader数据加载占用了大部分时间，各位大佬都是怎么解决的？ 第1页

如果可以的话，换硬盘。我猜提问者的数据是放在了机械硬盘上吧？其实你可以测试一下是不是单纯的数据读取阶段特别慢，因为根据我的经验这么长的读取时间已经不是transform造成的了。测试方法也很简单，你写dataset的时候，不要用opencv或者pillow对图片解析，直接用open()和read()读取每个图片的二进制数据（我没记错的话其实pillow中open图片时就是这么做的，你可以看看pillow中的源代码），dataloader返回的数据只是图片文件的二进制数据，这样测试dataloader，看看每一批时间是多久；可以读取每一批数据之后做一个delay，模拟训练模型花费时间。如果这样测试后发现每一批或者每隔几批后时间很久，最好是换硬盘吧。另外还有个简单的测试方法：移动/复制数据。比如把数据集的一部分或者找几个大文件，从硬盘复制到/dev/shm中，其实相当于从硬盘复制到内存中，计一下时间，就能得到传输速度，也可以根据你复制的文件量或者文件大小，换算出复制的数据集百分比，进而估计出一批数据读入内存所需的时间。

我在pytorch的数据读取上花了很大工夫做尝试，我说一下我的经历和经验。

我是做视频的，先说一下我的困境。假设，每一个batch包含16个视频，要在每个视频中取8帧，也就是送入模型的一个batch总共有128张图片。读取这128张图片有两种方法，一种方法是先读取整个mp4格式的视频，用openCV对每一帧解码然后返回要使用的帧，或用opencv里video.set方法跳跃到某帧对应的index只对该帧解码，当然前者运算量太大一般不用，使用后一种方式；另一种方法是提前将所有视频解码到图片，用的时候直接根据文件名读取图片就可以。这两种方法相比，第一种方法解码运算量大，因为从视频中解码帧比图片解码计算量大得多，查看CPU负载就可以看出差别来；第二种方法虽然快、计算量小，但是解码后的图片需要大量存储空间，保存到jpeg格式95%质量的情况下，像UCF101这样的视频数据集，图片数据是视频数据的6.5倍左右，Kinetics-400数据集解压到图片有1.5TB左右，如果想要保存成bmp格式或者numpy矩阵更不可能（一个图片的bmp格式与npy文件大小相同）。两种方法都有非常致命的缺点，只不过是牺牲计算还是牺牲存储空间的区别了。一般来说还是牺牲存储空间更合适，因为解码视频非常慢，即使数据从硬盘上的读取速度快，也会造成显卡有大量空余时间，并且CPU一直高负载运行也不利于服务器做其他任务（ssh和vim都卡）。另外，openCV从avi/mp4等格式里解析出来的图片数与用get()方法得到的视频帧数不一定相同（原因网上有解释），会导致帧检索错误。因此，我主要针对读取图片的方法尝试做改进。

1. 我一开始使用pre_fetch（可以在其他回答中看到类似的方法）办法，但是并没有用。
2. 我使用DALI做视频预处理，但是DALI的视频处理当时(2019.09)并不完善，我发现了bug……就反馈给开发者了，回复我说之后的版本会修复并做了标记。后来我也没用过DALI就没再看过了。
3. 我看到一些迷惑人的博客或帖子，以为pytorch的dataloader多进程是导致变慢的元凶，所以去研究了dataloader的源码。Dataloader工作时，首先主进程取一批数据在数据集中的index，然后其他进程（进程数由numworkers控制）中的一个进程对这一批数据进行读取和处理，处理完的数据放到dataqueue中，主进程从dataqueue中读取数据，然后再给一批数据的index。在dataloader开始使用时，都会出现数据加载时间很长的情况，这是因为第一次numworkers个进程都要产生2批数据。每次主进程产生数据index的代价都特别小，其实就是选择几个随机数字，网上有的地方说这一步导致数据不能并行，其实是错误的。但是，每个进程读取一批数据的时候是串行的，这一批中的数据一个一个读取。我就针对这一步做了改进尝试，在这个串行读取一批数据的代码里，源码有单线程限制语句，我就去掉了改成了多线程，我对这方面了解不够专业，不知道python里多进程加上多线程对不对，反正就试试吧。结果是，并没有提速效果。
4. 后来我专门测试了数据读取和数据转换的时间，发现经常会读取一张图片特别慢，1秒到2秒才能读到，几乎隔几张图片就会有这种卡顿，我试了PIL和openCV还有其他方法，都没办法提高速度。于是我按照文章开头的方法，做了硬盘读取速度的测试，发现是硬盘读取速度的限制造成了数据读取慢，而图片transform时间很短，也一直很稳定。
5. 针对硬盘读取数据，我想了一些办法。硬盘每次读取数据的时候，如果数据是连续的，那么读取速度会快很多，如果数据是很多小文件，读取速度就很慢，这一点平时复制文件的时候也会碰到。一批数据其实才16个视频，如果一次读取整个视频，就能减少读取次数，但是读取视频再解码运算量大。在一些框架（caffee、DALI）中会用到LMDB和HDF5文件格式，我可以借助这两种格式来把整个数据集所有图片保存到一个文件，从而减小读取次数。我看到一个帖子通过实验发现h5格式更快，所以就决定使用h5格式。但如果把每张图片用openCV读取后保存到numpy矩阵再存入h5文件，占用的存储空间太大（前边介绍过原因），所以就干脆保存每一个图片的二进制数据到numpy矩阵再存入h5文件，所有图片都读入，索引就是相对路径，这样做好了数据集文件。在使用的时候却发现了问题，h5文件不支持多进程读取，numworkers不能大于1，否则会报错；我又尝试了另一种格式，不是LMDB，而是npz格式，方法类似h5文件，但同样不支持多进程读取。
6. 单个h5和npz文件不支持，我就按视频拆分，每个视频保存一个h5文件，每个h5文件保存的是该视频每张图片的二进制数据，这样读取1个视频的时候一次性读取1个h5文件，取出里边的8张图片再解码就可以，相比直接读取视频找到帧再解码，这种解码图片的方法运算量低很多，当然对于同一个视频h5文件要比mp4/avi文件大很多，所以这种方法算是对运算量和占用空间的平衡吧。这种方法当时把训练时间降低了一点，效果不大。
7. 根据我的测试，h5文件读取数据的时候也并不是一次性读取整个文件的数据的，而是根据索引（或者说key）读取检索的数据，只不过一个h5文件是连续存储的。我还是想要一次性读取整个视频的内容，经过我测试，npy文件每次都是返回整个文件的数据，那我可以把每个视频所有图片的数据存在npy文件里。可是不同的图片数据长度不一样，一个npy文件只能保存维度相同的向量，所以我再一次牺牲空间，每个npy文件保存的矩阵里，矩阵的维度一个是视频帧数，另一个是这个视频所有图片中文件数据最长的长度，其他长度不足的图片后续补0。这样可以读入整个视频文件，直接根据矩阵索引找到对应的图片，读入图片未解码的数据后，通过PIL解码就可以还原图片了。当时我使用这个方法，把训练时间（间隔不是1个batch）从18秒左右减小到13秒左右。所以我一直使用这种方法了。
8. 但是我还是不满意。我想能不能直接从内存中读取npy文件呢？一次性加载所有数据集是不可能的，内存没那么大，那我可不可以一边把硬盘上的数据复制到内存里，一边读取内存的数据呢？可以先提前复制一些数据，给硬盘留一些时间余地。内存位置可以用/dev/shm，让dataloader到这里来读数据。我写程序做了试验，结果是这个方法不管用，因为从硬盘复制到/dev/shm太慢了。这一步也让我认识到了，数据加载慢只能怪硬盘。
9. 没有别的办法了，现在我就是使用7里边的npy文件数据集或者原始的视频文件，数据集不大的情况下会先把数据集复制到主硬盘SSD里，这个方法挺爽的。
10. 还没有结束，最近看到一篇文章：

“数据回波”，简单来说就是重复使用已经训练过还在内存中的数据。我还没有做试验。

更新：

之前实验了“数据回波”，太影响精度，提上来的一两个点又被抹平了，所以不建议使用。数据速度要加载快，归根到底还是要换固态或者内存，另外不要两个人同时从硬盘上读，可以分硬盘用。

又更新：

有些回答是针对数据预处理速度慢的，如果真的是这一步慢，也没必要使用DALI什么的，torchvision的Transforms中，几乎所有操作都支持PIL.Image和torch.Tensor两种输入，因此将图片读入后，先转换成Tensor（甚至torchvision有个read_image函数可以直接返回Tensor），然后放到GPU上，再用Transforms预处理，也是可以的。这个办法我已经用过了，在一个任务中我发现数据增强里的色彩增强最影响速度，其次是旋转，其他resize、flip都很快，因此用上述方法，在GPU上预处理数据。弊端是需要占用不少显存，我当时用的V100，还倒吃得消。为了能再节省显存，我把数据增强前的Tensor转换成了float16，送入网络之前再转换成float32，中间处理过程中，某些计算函数可能不支持float16，因此需要转换到float32再转换回float16。还需要注意的是，为了避免所有数据都只在GPU 0上预处理，要对不同的数据处理进程使用不同的GPU，比如可以用pid%n_gpu来指定。虽然麻烦点，哪怕中间还有float16和float32的来回转换，但是不怎么影响速度，数据处理的显存使用也能降低一些。亲测好用。