Pytorch采坑记录：每隔num_workers个iteration数据加载速度很慢

  最近在做某个视觉任务的模型训练，由于数据量比较少为了效果好一点，决定现在imagenet上pretrain一下骨干网络。但是在训练的时候遇到了一个问题：每隔num_workers个iteration数据加载都很慢，通过查找资料和代码搞清了这个问题。

背景

  设计了一个网络做目标检测，骨干网络是自己diy的因此没有pretrain的模型。而目标检测的数据集比较小，为了把模型训的好一点决定现把骨干网络搭一个分类头做个分类模型，在ImageNet上面pretrain一下。于是乎下载了imagenet的数据集，训练集总共120多万张图，验证集5万张，全部都按图片儿存储。自己写了一个dataset，然后打包成一个dataloader做训练的循环。为了加快数据读取速度，dataloader里面指定了num_workers=4。由于gpu显存比较充裕，设置了很大的batch_size=384.

问题

  在训练的时候发现，每当迭代次数能整除4的时候数据加载的速度都很慢，否则就比较快。改变num_workers为不同的值都会出现类似的情况，每当迭代次数能被num_workers整除时，数据加载都很慢，否则就比较快。只有当设置num_workers=0时即只通过主进程加载数据时，每一迭代的耗时才比较均匀。比如，当把num_workers设置为4时，抓下来的log:

train batch 1 load time 21.43259024620056s
train batch 2 load time 0.031423091888427734s
train batch 3 load time 0.004406452178955078
train batch 4 load time 0.004347562789916992
train batch 5 load time 18.13344931602478
train batch 6 load time 0.004399538040161133
train batch 7 load time 0.03353142738342285
train batch 8 load time 0.004467010498046875
train batch 9 load time 16.202253103256226
train batch 10 load time 0.8377358913421631
…

原因分析

  这个问题的原因其实比较简单：数据加载的速度太慢了，dataloader加载一个batch的时间远远大于gpu消费一个batch的时间，导致gpu大多时候都在等待。画个图来说明一下。
  首先构建了一个dataset，里面装着120多万张ImageNet的训练集图片，然后通过一个num_workers=4的dataloader从数据池里不停的打包batch。num_workers=4相当于开启了4个子进程分别独立的进行batch的打包，此时GPU在焦急的等待，他说你们4个进程小兄弟，任何一个打包完一个batch都立刻给我我好拿去前向传播。由于资源调度等细微差别，4个打包batch的进程速度基本相同但稍有区别。不妨假设某次打包，四个进程同时开始，进程1完成当前batch打包需要100ms，进程2完成当前batch打包需要98ms，进程3完成当前batch打包需要101ms，进程4完成当前batch打包需要102ms。

  那么从循环结果的角度来看，相当于需要等待98ms才能进行第一次循环，进程2辛辛苦苦打包好的一个batch被拿去消费了，然后他就默默的再去进行下一次打包了。而仅仅过了2ms进程1就又打包好了一个batch供给一次gpu循环，再过了仅仅1ms进行第三次循环，再过1ms进行第4次循环。此时共进行了4次循环，分别把4个进程辛苦了大约100ms才打包好的数据全用掉了。虽然进程2在第98ms就已经开始进行下一个batch的打包了，但是仍然需要再过大约100ms才能准备好下一次。而且由于资源调度等方面的随机原因，最早开始下一轮打包数据的进程，不见得在这一轮仍然是第一个完成打包的。
  单个进程打包一个batch需要大概100ms，4个进程同时工作结合上进程的调度相当于dataloader打包好一个batch的等效时间为:100/4=25ms，可是gpu消费数据的速度很快，比如只需要10ms。那么就会造成gpu在每个num_workers次循环都需要等待很长时间，看看时序图就清楚了。

  此时从循环时序的角度来看：每当迭代次数被子进程数整除时，迭代很慢因为被数据加载拖累；除此之外迭代都很快，因为gpu刚消费完一个batch另外一个batch也已经准备好了，gpu直接拿去再次消费没有等待时间。

解决办法

  根据上面讨论的原因，不难得出以下几种可能的解决办法，最根本的逻辑就在于：要让多个子进程并行情况下的等效batch打包时间，恰好等于gpu的消费batch的时间，此时既不会出现dataloader已经打包好了batch等待gpu使用，也不会出现gpu算力已经空闲等待dataloader打包batch ，可以充分的将整个pipeline的资源全部利用起来。

增加num_workers数值

  dataloader的打包batch等效时间理论上等于：单个进程打包时间/进程数量。因此增加num_workers数值，增加进程数可以缩小数据打包和gpu消费数据之间的时间差。但是需要注意的是进程之间的调度也是有时间消耗的，因此并不是进程数越多打包速度越快。当打包时间和gpu消耗batch时间相差不大时，可以考虑通过这种方法优化。

减小数据读取时间

  进程打包时间很长的一个重要原因是，我在做ImageNet预训练时原始数据是一张张的图片儿，batch size设置为384是一个比较大的数，单个进程打包一个batch需要进行很多张图片文件的打开、读取和关闭操作，IO时间远远大于数据读取和处理时间，因此可以考虑减小数据的IO操作次数提高数据读取速度，例如类似于tensorflow把数据转成tfrecord这种大尺寸的二进制文件等。
  有关这个部分，可以看一下我写的另外一篇讲dali的文章，完美的解决了数据读取的时间问题，gpu利用率可以达到几乎100%。

  除了上面的2种，应该还有其他办法解决这个问题，这里就不多罗列了。水平有限，欢迎讨论。