降低訓練精度提高batch大小，512*512的圖片在8G的顯存上batch只能為1，2疏之，使用了apex后可以增大到10。只支持pytorch暇咆。
官網(wǎng)地址：NVIDIA/apex
官方說支持Ubuntu系統(tǒng)锋爪，Windows只是實驗性的，但是我在win10上也安裝成功了爸业。
前提
- 已安裝好CUDA和CUDNN其骄，且版本適配。
2.在安裝前先檢查一下扯旷，電腦的cuda版本和pytorch內(nèi)的cuda版本是否一樣拯爽，不一樣的話就把低版本的進行升級。
# 查看電腦的cuda版本钧忽、
>> nvcc -V
# pytorch內(nèi)的cuda版本
import torch
torch.version.cuda
安裝
按照官網(wǎng)的命令輸入即可
$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
如果報錯出現(xiàn)什么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()
當中斷后繼續(xù)訓練時
# 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
···
用戶指定數(shù)據(jù)格式
amp.initialize(net, opt, opt_level=“O1”)
其中的opt-level參數(shù)是用戶指定采用何種數(shù)據(jù)格式做訓練的。
- O0：純FP32訓練洲拇，可以作為accuracy的baseline奈揍；
- O1：混合精度訓練（推薦使用）曲尸，根據(jù)黑白名單自動決定使用FP16（GEMM, 卷積）還是FP32（Softmax）進行計算。
- O2：“幾乎FP16”混合精度訓練男翰，不存在黑白名單另患，除了Batch norm，幾乎都是用FP16計算蛾绎。
- O3：純FP16訓練昆箕，很不穩(wěn)定，但是可以作為speed的baseline租冠；
溢出問題
相乘的操作鹏倘，16位*16位=32位，精度不會有大的影響顽爹，但16位+32位由于后面32位的數(shù)太小纤泵，會進行舍入只得到一個16位數(shù)，影響較大镜粤，所以當涉及到“加法”時最好恢復到32位精度捏题，比如 sigmoid,softmax。
我們僅需在模型定義中肉渴，在構(gòu)造函數(shù)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指明后面的函數(shù)需要使用float類型。注意第二實參是string類型
和register_float_function相似的注冊函數(shù)還有
- amp.register_half_function(module, function_name)
- amp.register_float_function(module, function_name)
- amp.register_promote_function(module, function_name)
你必須在使用amp.initialize之前使用注冊函數(shù)黄虱，所以最好的位置就放在模型的構(gòu)造函數(shù)中
