Pytorch显存管理
在模型训练过程中,显存的管理是需要注意的环节,尽量用满显存在一定程度上可以缩短训练时间,但过于贪心往往会遇到CUDA out of memory问题,钞能力不足的情况下,可以尝试一下手段来减少资源需求。
减少batch_size:这是最简单有效的方法,减少一批读取的数据,从而减少显存的使用量,但如果每个批次都需要重新读取新的数据而不是固定死的,又没办法一次性读取完数据集(比较少见这样的需求,但我遇到了),就得尝试其他的方法了。
勤使用torch.cuda.empty_cache():虽然python内部也有资源管理机制,但有时候还是笨笨的,不能及时释放资源,此时就需要我们调用该函数手动释放。
减少线程使用num_workers:在dataloader中将会传入该参数,可以减少该参数值从而减少内存使用。
使用混合精度:计算机中数据存储的格式有许多种,此处举个简单例子如short int和int,之间就差距很大,pytorch中有torch.cuda.amp模块可以来实现混合精度训练,需要注意的是,该方法会存在精度损失,从而影响最终迭代结果,但因为一般情况下损失微乎其微,所以基本不会对结果产生影响,在这方面有需求的使用者慎重,下面是示例代码:
最后如果以上方法都无法满足,就需要从模型入手,深度可分离卷积是一个很好地降低模型参数的方法,当然优化数据也是可以的,相信能做到这些的同学,也不需要我过多说明。
最后希望大家都有钞能力能免受这种琐碎之事的困扰。