前言
最近的工作中，用到了Pytorch框架训练医学图像分割模型。精心设计的模型经常会因为显存不足而失败。减小模型训练过程中对显存的占用，可能我们能想到最简单的方法就是减小batchsize，减少卷积核数量，裁剪输入图像的大小等。但是，以上方法可能会影响模型性能。经过多次尝试，总结了几种尽可能不改变模型结构，不影响模型性能且能够节省显存的训练方法。
清理GPU缓存
Pytorch提供了torch.cuda.empty_cache()函数用来清理GPU的缓存数据，在实际的模型训练中，我们可以将此函数与异常处理相结合，一旦发生显存溢出，异常处理机制可以捕获异常信息，清理GPU缓存，以保证模型正常训练。总体实现代码如下：
...
try:
output_data = model(input_data)
loss = calcloss(output_data,input_label)
...
except RuntimeError as exception:
if "out of memory" in str(exception):
print("GPU 显存不足！")
if hasattr(torch.cuda,"empty_cache"):
torch.cuda.empty_cache()
else:
raise exception
...
此方法能够起作用的前提是模型自身占用显存较小，如果模型过大，很容易丢失大量训练数据，影响模型性能。
改进反向传播中activation的存储方式
Tianqi Chen的论文中提出的方法中，通过将CNN反向传播过程中，每个节点存储activation的过程改进为每隔sqrt(N)个节点存储一个activation的方法，将原有显存占用从O(N)降至O(sqrt(N)),非常有效地减小了显存的占用。
已经有大佬用pytorch实现了该方法，一个简单的实现方式是将github中相关代码和模型放入同一目录，将所有nn.Sequential()替换为SublinerSequential()方法实现。SublinerSequential()存储的层数越多，节省显存的效果越明显。具体示例代码如下所示：
from memonger import SublinearSequential
self.layer = SublinearSequential(nn.Conv3d(...),
nn.BatchNorm(...),
nn.Relu(...),
...
)
对于一个结构较为复杂的模型，可以利用回调函数的方法，将结构相同，参数不同的模块构造为一个函数实现。该项目也提供了一个resnet的实现方式，以提供参考。
NVIDIA apex 方法
NVIDIA提供了apex的方法，在Pytorch中，通过16位混合浮点运算，不仅能够节省显存，也能加快模型训练速度。
apex的安装方式如下：
$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
如果第3条指令执行报错，可以用以下指令代替，基本不影响模型训练。
$ pip install -v --disable-pip-version-check --no-cache-dir ./
apex安装成功后，对模型的训练代码做以下改进，以实现加速训练和节省显存的目地。
from apex import amp
model,optimizer = amp.initalize(model,optimizer,opt_level="O1")
output_data = model(input_data)
loss = loss(output_data,input_label)
with amp.scale_loss(loss,optimizer) as scaled_loss:
scaled_loss.backward()
...
apex节省显存的效果非常明显，有关apex更多的技巧和介绍可以关注apex的官方github仓库:
https://github.com/NVIDIA/apex
总结
将以上所有方法结合一起使用，在节省显存方面，效果非常明显。当然,更好的办法是换一台显存更大的电脑（^_^）
参考：
[1] https://www.zhihu.com/question/274635237
[2] https://github.com/Lyken17/pytorch-memonger
[3] https://github.com/NVIDIA/apex
