dataloader合理设置num_works和batchsize 避免爆内存

dataloader合理设置num_works和batchsize,避免爆内存


个人总结,禁止以任何形式的转载!!

1.关乎内存

1)dataloader会开启num_works个进程,如图所示:(这里设置的是6)

在这里插入图片描述

ps -eH

在这里插入图片描述

2)以图片为例,每个worker通过主进程获得自己需要采集的图片id(id的顺序由sampler或shuffle得到)。然后每个worker开始采集一个batch的数据,存放在内存。(因此增大num_workers,内存占用也会增加。因为每个worker都需要缓存一个batch的数据),查看内存使用情况:

free -mh

在这里插入图片描述
因此,设置 num_workers X batchsize 过大的话,内存会不够用导致程序中止掉:
DataLoader worker (pid ****) is killed by signal: Killed.

3)在第一个worker数据采集完成后,会卡在这里,等着主进程把该batch取走,然后采集下一个batch。主进程运算完成,从第二个worker里采集第二个batch,以此类推。(->1 ->2 ->3 ->4-> 5 ->6 ->1 ->2 ->3-> 4-> 5 ->6…)

4)主进程采集完最后一个worker的batch。此时需要回去采集第一个worker产生的第二个batch。如果该worker此时没有采集完,主线程会卡在这里等。(这种情况出现在,num_works数量少或者batchsize 比较小,显卡很快就计算完了,CPU对GPU供不应求。)

2.关乎显存

GPU的内存简称显存,这里存放着送入显卡的数据,以及一些参数和中间变量。影响最大的就是参数量和一个batch的数据量,当batchsize过大时,会爆显存,太小时显卡的性能没有真正展现出来。查看显卡工作状态:

nvidia-smi

在这里插入图片描述
我这里batchsize=1,所以显卡运算单元利用率并不高,而且显存也还剩余特别多。

注意:显存占用不是和batchsize简单成正比,模型自身的参数及其延伸出来的数据也要占据显存。

3.总结

想要加快训练速度,就需要利用num_works和batchsize协调好CPU和GPU的资源,让GPU单元利用率尽量高,同时内存的也不会溢出。其中batchsize对模型性能有所影响,故需要优先考虑batchsize的大小。

在有的代码中你可能会看见pin_memory。pin_memory就是锁页内存,创建DataLoader时,设置pin_memory=True,则意味着生成的Tensor数据最开始是属于内存中的锁页内存,这样将内存的Tensor转义到GPU的显存就会更快一些。主机中的内存,有两种存在方式,一是锁页,二是不锁页,锁页内存存放的内容在任何情况下都不会与主机的虚拟内存(Swap)进行交换(注:虚拟内存就是硬盘),而不锁页内存在主机内存(Mem)不足时,数据会存放在虚拟内存中。当计算机的内存充足的时候,可以设置pin_memory=True。当系统卡住,或者交换内存使用过多的时候,设置pin_memory=False。

个人总结,禁止以任何形式的转载!!
抛砖引玉,欢迎各位交流和指正!!

  • 49
    点赞
  • 101
    收藏
    觉得还不错? 一键收藏
  • 15
    评论
评论 15
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值