Nvidia Apex安装与简单使用

降低训练精度提高batch大小，512*512的图片在8G的显存上batch只能为1，2，使用了apex后可以增大到10。只支持pytorch。

官网地址：NVIDIA/apex
官方说支持Ubuntu系统，Windows只是实验性的，但是我在win10上也安装成功了。

前提

已安装好CUDA和CUDNN，且版本适配。

2.在安装前先检查一下，电脑的cuda版本和pytorch内的cuda版本是否一样，不一样的话就把低版本的进行升级。

# 查看电脑的cuda版本、
>> nvcc -V

# pytorch内的cuda版本
import torch
torch.version.cuda

安装

按照官网的命令输入即可

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

如果报错出现什么pip 137..我不记得了
在保证cuda版本一致的前提下进入apex文件夹，使用命令：python setup.py install 即可安装成功

验证安装成功

from apex import amp

没有报错就是成功了。

Apex的使用

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()

当中断后继续训练时

# 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；

溢出问题

相乘的操作，16位*16位=32位，精度不会有大的影响，但16位+32位由于后面32位的数太小，会进行舍入只得到一个16位数，影响较大，所以当涉及到“加法”时最好恢复到32位精度，比如 sigmoid,softmax。
我们仅需在模型定义中，在构造函数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之前使用注册函数，所以最好的位置就放在模型的构造函数中

参考链接：
英伟达（NVIDIA）训练深度学习模型神器APEX使用指南

作者：zelda2333
链接：https://www.jianshu.com/p/f70ca1d0d91b
来源：简书
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。