降低训练精度提高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之前使用注册函数，所以最好的位置就放在模型的构造函数中
