英伟达（NVIDIA）训练深度学习模型神器APEX使用指南
你是否苦闷于教研室卡不多，卡显存不大，很多模型没法跑，是否发愁不能用很大的batch size导致loss没法降低。如果你使用的是PyTorch，恭喜你，你完全可以使用APEX从中解脱出来。
APEX是什么
APEX是英伟达开源的，完美支持PyTorch框架，用于改变数据格式来减小模型显存占用的工具。其中最有价值的是amp（Automatic Mixed Precision），将模型的大部分操作都用Float16数据类型测试，一些特别操作仍然使用Float32。并且用户仅仅通过三行代码即可完美将自己的训练代码迁移到该模型。实验证明，使用Float16作为大部分操作的数据类型，并没有降低参数，在一些实验中，反而由于可以增大Batch size，带来精度上的提升，以及训练速度上的提升。
接下来开始配置APEX
APEX的配置
简单。前提是你安装好了CUDA和CUDNN，以及你的系统是Ubuntu系统。
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
注意安装可能会报错
Cuda extensions are being compiled with a version of Cuda that does " + not match the version used to compile Pytorch binaries. " + Pytorch binaries were compiled with Cuda
原因是你的CUDA版本和Pytorch版本对不上，尽管你能使用支持GPUI的Pytorch。现在查看一下版本。
查看CUDA版本，使用
nvcc -V
查看torch的cuda版本
import torch torch.version.cuda
对比一下，如果cuda版本低就升级cuda，如果Pytorch版本低就升级Pytorch。之后重复上面安装过程就OK了。 安装之后，clone下来的apex文件夹就可以删除了。 查看能否正确导入apex
from apex import amp
net = xxxNet()
net.cuda()
net.train()
params_low_lr = []
params_high_lr = []
for n, p in net.named_parameters():
if 'encoder' in n:
params_low_lr.append(p)
else:
params_high_lr.append(p)
opt = Adam([{'params': params_low_lr, 'lr': 5e-5},
{'params': params_high_lr, 'lr': 1e-4}], weight_decay=settings.WEIGHT_DECAY)
net, opt = amp.initialize(net, opt, opt_level="O1") # 这里要添加这句代码
...
...
loss = ...
if not torch.isnan(loss):
opt.zero_grad()
with amp.scale_loss(loss, opt) as scaled_loss:
scaled_loss.backward() # loss要这么用
opt.step()
然后你舅舅可以惊喜的发现，同样的网络和Batch size，显存用的少了。尽情增大Batch size吧。
当你需要使用resume的方式训练的时候。按照以下方式
# Initialization
opt_level = 'O1'
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
# Train your model
...
# Save checkpoint
checkpoint = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'amp': amp.state_dict()
}
torch.save(checkpoint, 'amp_checkpoint.pt')
...
# Restore
model = ...
optimizer = ...
checkpoint = torch.load('amp_checkpoint.pt')
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model.load_state_dict(checkpoint['model']) # 注意，load模型需要在amp.initialize之后！！！
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])
# Continue training
...
用户指定数据格式
amp.initialize(net, opt, opt_level=“O1”)
其中的opt-level参数是用户指定采用何种数据格式做训练的。
O0：纯FP32训练，可以作为accuracy的baseline；O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算。O2：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。O3：纯FP16训练，很不稳定，但是可以作为speed的baseline；
溢出问题（必看）
因为Float16保存数据位数少了，能保存数据的上限和下限的绝对值也小了。如果我们在处理分割类问题，需要用到一些涉及到求和的操作，如sigmoid，softmax，这些操作都涉及到求和。分割问题特征图都很大，求个sigmoid可能会导致数据溢出，得到错误的结果。所以针对这些操作，仍然使用float32作为数据格式。那么如何基于上面的代码修改呢。
我们仅需在模型定义中，在构造函数__init__中的某一个位置。加上下面这段：
from apex import amp
class xxxNet(Module):
def __init__(using_map=False)
...
...
if using_amp:
amp.register_float_function(torch, 'sigmoid')
amp.register_float_function(torch, 'softmax')
用register_float_function指明后面的函数需要使用float类型。注意第二实参是string类型 和register_float_function相似的注册函数还有
amp.register_half_function(module, function_name)amp.register_float_function(module, function_name)amp.register_promote_function(module, function_name)
你必须在使用amp.initialize之前使用注册函数，所以最好的位置就放在模型的构造函数中
关于apex分布式点击下一篇
英伟达APEX，多GPU分布式训练，同步Batchnorm，自动混合精度训练法宝指南
参考
Pytorch 安装 APEX 疑难杂症解决方案 【PyTorch】唯快不破：基于Apex的混合精度加速