pytorch计算模型的显存占用率和节省内存技巧

计算模型的显存占用率

• 如何计算模型以及中间变量的显存占用大小可以参考此文。
• 如何在Pytorch中精细化利用显存，牺牲计算速度减少显存用量，将计算过程分为两半，先计算一半模型的结果，保存中间结果再计算后面一半的模型。code
• 再次浅谈Pytorch中的显存利用问题(附完善显存跟踪代码)，借用Pytorch-Memory-Utils这个工具来检测我们在训练过程中关于显存的变化情况，分析出我们如何正确释放多余的显存。

节省内存的技巧

• 减小Batch-size；
• 减小图片尺寸；
• 将计算过程分为多步，保存中间结果再计算后面一半的模型，参考此文；
• 使用软batch训练model，即通过修改优化算法，是网络进行N次前向传播，但只进行一次反向传播；
• 使用pooling，减小特征图的size；
• 减少全连接层的使用。relu(inplace=true)，inplace_abn使用半精度float16；
• optimizer的变换使用，理论上，sgd<momentum<adam，可以从计算公式中看出有额外的中间变量；
• Depthwise Convolution；
• 尽可能使用inplace操作， 比如relu 可以使用 inplace=True；
• 不要一次性把数据加载进来，使用 pytorch 的 Dataset api；
  大致程序如下：

  from torch.utils.data import DataLoader
  data_train_loader = DataLoader(data_train, batch_size=64, shuffle=True)
  for batch_idx, samples inenumerate(data_train_loader):
  net(samples)

优化pytorch代码

• https://www.sagivtech.com/2017/09/19/optimizing-pytorch-training-code/
• Memory-Efficient Implementation of DenseNets