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