Pytorch节省显存、加速训练的小技巧
Pytorch节省显存、加速训练的小技巧
2020.7.22
张文政
GPU是深度学习模型的关键，对GPU来说，最重要的两个指标就是显存与算力。如何在代码层面最大程度地节省显存、在算力固定的情况下更快地训练，是值得学习地。当然，APEX和多GPU训练可以显著地提升这两点，本文只介绍除了APEX和多GPU之外的其他的技巧。
使用4层编码，4层解码，8头注意力的transformer，优化器使用Adagrad，在生成任务上进行测试。初始时训练batchsize为16，占用显存12000MB；测试batchsize为16，占用显存15000MB。
文章目录
- Pytorch节省显存、加速训练的小技巧
-
- 显示显存与GPU利用率
- 到底什么在占用显存？
- 技巧1：inplace=True
- 技巧2：with torch.no_grad():
- 技巧3：forward中的变量命名
- 技巧4：Dataloader数据读取
- 技巧5：gradient accumulation
- 参考
显示显存与GPU利用率
- nvidia-smi
- 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
[参考]
一些激活函数与Dropout有一个参数"inplace"，默认设置为False，当设置为True时，我们在通过relu()计算时的得到的新值不会占用新的空间而是直接覆盖原来的值，这也就是为什么当inplace参数设置为True时可以节省一部分内存的缘故。但在某些需要原先的值的情况下，就不可设置inplace。
此操作相当于针对显存占用第二点（模型中间变量）的优化。
技巧2：with torch.no_grad():
[参考讨论]
对于只需要forward而不需要backward的过程（validation和test），使用torch.no_grad做上下文管理器（注意要在model.eval()之后），可以让测试时batchsize扩大近十倍，而且也可以加速测试过程。此操作相当于针对显存占用第二点（因为直接没有backward了）和第三点进行优化。
model.eval()
with torch.no_grad():
pass
技巧3：forward中的变量命名
在研究pytorch官方架构和大神的代码后可发现大部分的forward都是以x=self.conv(x)的形式，很少引入新的变量，所以启发两点以减少显存占用（1）把不需要的变量都用x代替，（2）变量用完之后马上用del删除（此操作慎用，清除显存的同时使得backProp速度变慢）。此操作相当于针对第二点（模型中间变量）进行优化。
技巧4：Dataloader数据读取
一定要使用pytorch的Dataloader来读取数据。按照以下方式来设置：
loader = data.Dataloader(PYTORCH_DATASET, num_works=CPU_COUNT,
pin_memory=True, drop_last=True)
第一个参数是用pytorch制作的TensorDataset，第二个参数是CPU的数量（默认为0，在真正训练时建议调整），第三个参数默认为False，用来控制是否把数据先加载到缓存再加载到GPU，建议设置为True，第四个参数用于扔掉最后一个batch，使得训练更为稳定。
将pin_memory开启后，在通过dataloader读取数据后将数据to进GPU时把non_blocking设置为True，可以大幅度加快数据计算的速度。
for input_tensor in loader:
input_tensor.to(gpu, non_blocking=True)
model.forward(input_tensor)
技巧5：gradient accumulation
[参考]
梯度积累通过累计梯度来解决本地显存不足的问题，即不在每个batch都更新模型参数，而是每经过accumulation steps步后，更新一次模型参数。相当于针对第三点（n步才更新一次参数）来进行优化。且由于参数更新的梯度计算是算力消耗的一部分，故梯度累计还可以一定程度上加快训练速度。
loss = model(input_tensor)
loss.backward()
if batch_idx % accumulate_steps == 0:
optim.step()
optim.zero_grad()
相当于一个epoch的步数（step）变少了（一个step相当于参数更新一次），但单个step的计算时间变长了（略小于n倍的原来时间）
参考
[1]https://www.zhihu.com/question/274635237
[2]https://blog.csdn.net/qq_28660035/article/details/80688427?utm_medium=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromMachineLearnPai2-1.nonecase&depth_1-utm_source=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromMachineLearnPai2-1.nonecase
[3]https://discuss.pytorch.org/t/model-eval-vs-with-torch-no-grad/19615
[4]https://oldpan.me/archives/how-to-use-memory-pytorch
[5]https://zhuanlan.zhihu.com/p/31558973