欢迎光临散文网 会员登陆 & 注册

Pytorch显存管理

2023-08-14 10:25 作者:傑君  | 我要投稿

在模型训练过程中,显存的管理是需要注意的环节,尽量用满显存在一定程度上可以缩短训练时间,但过于贪心往往会遇到CUDA out of memory问题,钞能力不足的情况下,可以尝试一下手段来减少资源需求。

  1. 减少batch_size:这是最简单有效的方法,减少一批读取的数据,从而减少显存的使用量,但如果每个批次都需要重新读取新的数据而不是固定死的,又没办法一次性读取完数据集(比较少见这样的需求,但我遇到了),就得尝试其他的方法了。

  2. 勤使用torch.cuda.empty_cache():虽然python内部也有资源管理机制,但有时候还是笨笨的,不能及时释放资源,此时就需要我们调用该函数手动释放。

  3. 减少线程使用num_workers:在dataloader中将会传入该参数,可以减少该参数值从而减少内存使用。

  4. 使用混合精度:计算机中数据存储的格式有许多种,此处举个简单例子如short int和int,之间就差距很大,pytorch中有torch.cuda.amp模块可以来实现混合精度训练,需要注意的是,该方法会存在精度损失,从而影响最终迭代结果,但因为一般情况下损失微乎其微,所以基本不会对结果产生影响,在这方面有需求的使用者慎重,下面是示例代码:

最后如果以上方法都无法满足,就需要从模型入手,深度可分离卷积是一个很好地降低模型参数的方法,当然优化数据也是可以的,相信能做到这些的同学,也不需要我过多说明。

最后希望大家都有钞能力能免受这种琐碎之事的困扰。

Pytorch显存管理的评论 (共 条)

分享到微博请遵守国家法律