pytorch 关于显存增长原因以及显存占用优化

ps:师傅希望想tensorflow一样设置 GPU占用(tf.GPUOptions),但事实情况是好像pytorch根本就没有相关函数来设置(如8G显存我只用2G,来跑是否可行)

1.网上很多攻略关于pytorch 训练时显存持续增长的问题

其实这个问题主要出现在0.4之前的版本,在累加loss时用的是loss_sum += loss,而非loss_sum += loss.data[0].因为loss其实是tensor,会持续不断的放进显存里面的.

0.4之后的版本已经用loss.item() 替代loss.data[0],获取的直接是数字.

2.占用优化.

写了一些,发现有人以及总结的十分全面了,我就再稍微补充下写.最原始参考:https://www.zhihu.com/question/274635237

博士就是吊.

我从另外一个地方搬运:https://v2ex.com/t/532315

在不修改网络结构的情况下, 有如下操作:

1. 同意 @Jiaming, 尽可能使用 inplace 操作, 比如 relu 可以使用 inplace=True

2. 进一步,比如 ResNet 和 DenseNet 可以将 batchnorm 和 relu 打包成 inplace,在 bp 时再重新计算。使用到了 pytorch 新的 checkpoint 特性,有以下两个代码。由于需要重新计算 bn 后的结果,所以会慢一些。
gpleiss/efficient_densenet_pytorch ( https://github.com/gpleiss/efficient_densenet_pytorch)
mapillary/inplace_abn ( https://github.com/mapillary/inplace_abn)

3. 每次循环结束时 删除 loss,可以节约很少显存,但聊胜于无。可见如下 issue

Tensor to Variable and memory freeing best practices ( https://discuss.pytorch.org/t/tensor-to-variable-and-memory-freeing-best-practices/6000/2)


4. 使用 float16 精度混合计算。我用过 
@NVIDIA 英伟达
apex,很好用,可以节约将近 50%的显存,但是要小心一些不安全的操作如 mean 和 sum,溢出 fp16。

NVIDIA/apex ( https://github.com/NVIDIA/apex)


5. 对于不需要 bp 的 forward,如 validation 请使用 torch.no_grad , 注意 model.eval() 不等于 torch.no_grad() 请看如下讨论。

'model.eval()' vs 'with torch.no_grad()' ( https://discuss.pytorch.org/t/model-eval-vs-with-torch-no-grad/19615)


6. torch.cuda.empty_cache() 这是 del 的进阶版,使用 nvidia-smi 会发现显存有明显的变化。但是训练时最大的显存占用似乎没变。大家可以试试。

How can we release GPU memory cache? ( https://discuss.pytorch.org/t/how-can-we-release-gpu-memory-cache/14530)


7.另外,会影响精度的骚操作还有:

把一个 batchsize=64 分为两个 32 的 batch,两次 forward 以后,backward 一次。但会影响 batchnorm 等和 batchsize 相关的层。


相关链接:

老外写的提高 pytorch 效率的方法,包含 data prefetch 等

Optimizing PyTorch training code ( https://www.sagivtech.com/2017/09/19/optimizing-pytorch-training-code/)

我对其中几点做下补充:

6.其实https://oldpan.me/archives/pytorch-gpu-memory-usage-track中所写其实Pytorch已经可以自动回收我们“不用的”显存,类似于python的引用机制,当某一内存内的数据不再有任何变量引用时,这部分的内存便会被释放.Pytorch的开发者也对此进行说明了,这部分释放后的显存可以用,只不过不在Nvidia-smi中显示罢了。 torch.cuda.empty_cache()其实没啥用处.

7.其实这个功能已经在Pytorch-0.4.0出现了可以将一个计算过程分成两半,也就是如果一个模型需要占用的显存太大了,我们就可以先计算一半,保存后一半需要的中间结果,然后再计算后一半。具体参考https://blog.csdn.net/TTdreamloong/article/details/84886621,使用了checkpoint包.显存实在不够的同学可以试试.

猜你喜欢

转载自blog.csdn.net/qq_36401512/article/details/96163940