Pytorch dataloader加载数据很慢

现象

写了个模型,训练速度很慢,发现大部分时间都花在了加载数据的过程中,训练时间反而不多

主机内存占用(已提交相当于总申请内存,其中包含交换分区的大小,其后面的数字为总可申请大小,其会根据申请情况动态扩容):

解决

num_workers

DataLoader中设置参数 num_workers=3 效果:

主机内存占用:

测试过程中貌似有一次内存申请越来越多,进而导致磁盘被占满(交换分区)

该参数默认情况下为0

该参数用于指定加载数据使用的线程数

当其为0时,其顺序执行以下步骤:

  1. dataloader获取本轮训练需要的数据索引
  2. 主线程拿到索引后再调用 torch.utils.data.Dataset 的 getitem 方法逐个获取数据
  3. 获取完数据后,开始执行模型训练
  4. 一轮训练结束之后,再次返回到步骤1获取数据,如此往复,完成训练

当其>0时,真实的过程我还是不清楚,以下是我做的一个测试结果:

dataloader会在初次迭代时根据 batch_sampler 获取多个batch的索引(具体为 prefetch_factor*num_workers 个),并将他们传给workers。每个worker都将单独加载一个batch(例如batch size为16,则每个worker会加载16条数据)到主机内存。

真正训练时,会去主机内存中找到所需的batch(如果用到GPU,则会再次从主机内存复制到显存)。与此同时,也会再次调用workers去加载一批新的数据

若没找到,则会再次调用workers去加载(这一句来源于网络)

实测结果显示,貌似只有在首次迭代dataloader时会一次加载多批数据,后续每次都是只用一个worker加载一批数据。我想这样应该也是合理的,如果每次都加载多批数据,消耗速度小于加载速度,最终内存会爆掉

经测试:

在我的某个模型中,每增加一个worker(默认prefetch_factor=2)内存消耗增加5.8G(绝大部分为交换分区,所以一旦worker多了之后就会大量占用磁盘空间,得注意)

可能是因为数据集较小(765M的纯文本数据),修改batchsize并不增加内存消耗,修改prefetch_factor对内存影响貌似也不是很大

以下为我某次测试数据(仅供参考,实际得依照具体模型和数据规模来看)

4worker 20次训练耗时16s 平均耗时0.8s
3…………………………………18s………………0.9s
2…………………………………28s………………1.4s
1…………………………………50s………………2.5s
0…………………………………50s………………2.5s

参考:

https://pytorch.org/docs/master/data.html

https://stackoverflow.com/questions/53998282/how-does-the-number-of-workers-parameter-in-pytorch-dataloader-actually-work

看到一个简洁明了的:https://www.cnblogs.com/h694879357/p/16055835.html

加了该参数程序就不动了?

我在windows10下使用jupyter,加入该参数后程序就卡在那了。

解决方法,导出为python文件,并且:1. 将方法体放到 main 方法中, 2. 使用命令行或其他ide执行

原因是设置 num_workers>0 需要多线程支持,jupyter notebook在这方面有很大的问题,并且还需要在 if __name__ == '__main__' 代码块中运行

参考:https://github.com/pytorch/pytorch/issues/51344

pin_memory

dataloader参数,默认为false,该值设置为true的结果(貌似影响没那么大)

该参数将指定一片固定的物理内存用于存放训练过程中需要的batch数据,遍历dataloader时,就会从该区域取得数据。由于内存地址是固定的,其寻址速度应该会快很多。

当使用GPU进行模型训练时,会需要将主机内存的batch数据拷贝到显存,该参数为true可加快这一拷贝动作。

再者,上述拷贝动作默认是主线程阻塞的,可以使用 tensor.to(device, non_blocking=True) 参数指定为异步复制

注:我也没搞清楚这一块的原理,我对文档中的 “This (non_blocking=True) can be used to overlap data transfers with computation.” 有疑问,数据的加载和计算同时进行是什么意思?加载的不是本次训练所需要的数据吗?

这篇文章貌似回答了这个问题:https://blog.csdn.net/silencez_w/article/details/121848437

意思是说,gpu计算时,只会用到数据的一部分,故可以将数据拆分成多个部分同时拷贝,这样当需要用到的数据拷贝完成后gpu就能立即进行计算,而其他部分数据仍然在继续拷贝

后记

以上方法我觉得都不好,要么可能会爆内存,要么提升不大。

为什么他不提供一个预加载的方法:在训练过程中异步预加载后一批数据即可。我看网上确实有很多人做了这样的实现

我有一个nlp的模型,数据不大(七百多兆),但长度大。导致每轮数据的加载都很慢,我想,如果能一次将所有的数据都加载进内存或显存,应该能快不少。

另外,我觉得可以用一个阻塞队列,在线程A中使用常规的dataloader向里面push数据,线程B读取队列头数据训练网络,由于两者是异步的,则开销基本就是训练部分的时间开销了(如果单次训练时常比单批数据加载时间更短)

后面对一次优化训练过程做了记录,见: https://blog.woyou.cool/post/%e8%ae%b0%e4%b8%80%e6%ac%a1%e6%a8%a1%e5%9e%8b%e8%ae%ad%e7%bb%83%e9%80%9f%e5%ba%a6%e4%bc%98%e5%8c%96%e7%9a%84%e8%bf%87%e7%a8%8b/

Leave a Comment