Pytorch节省显存、加速训练的小技巧

Pytorch节省显存、加速训练的小技巧

2020.7.22

张文政

GPU是深度学习模型的关键，对GPU来说，最重要的两个指标就是显存与算力。如何在代码层面最大程度地节省显存、在算力固定的情况下更快地训练，是值得学习地。当然，APEX和多GPU训练可以显著地提升这两点，本文只介绍除了APEX和多GPU之外的其他的技巧。

使用4层编码，4层解码，8头注意力的transformer，优化器使用Adagrad，在生成任务上进行测试。初始时训练batchsize为16，占用显存12000MB；测试batchsize为16，占用显存15000MB。

显示显存与GPU利用率

1. nvidia-smi
2. gpustat(推荐)

pip install gpustat
watch --color -n1 gpustat -cpu

到底什么在占用显存？

[参考]

输入的数据占用空间其实并不大，比如一个(256, 3, 100, 100)的Tensor（相当于batchsize=256的100*100的三通道图片。）只占用31M显存。

实际上，占用显存的大头在于：1. 动辄上千万的模型参数；2. 模型中间变量；3. 优化器中间参数[见1.2.2]。

第一点模型参数不必介绍；第二点，中间变量指每个语句的输出。而在backward时，这部分中间变量会翻倍（因为需要保留原中间值）。第三点，优化器在梯度下降时，模型参数在更新时会产生保存中间变量，也就是模型的params在这时翻倍。

技巧1：inplace=True

[