PyTorch 源碼解讀之 torch.cuda.amp: 自動(dòng)混合精度詳解
點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”
設(shè)為星標(biāo)，干貨直達(dá)！
AI編輯：我是小將
本文作者：OpenMMLab?@202011
https://zhuanlan.zhihu.com/p/348554267
本文已由原作者授權(quán)轉(zhuǎn)載
Nvidia 在 Volta 架構(gòu)中引入 Tensor Core 單元，來支持 FP32 和 FP16 混合精度計(jì)算。也在 2018 年提出一個(gè) PyTorch 拓展 apex，來支持模型參數(shù)自動(dòng)混合精度訓(xùn)練。自動(dòng)混合精度（Automatic Mixed Precision, AMP)訓(xùn)練，是在訓(xùn)練一個(gè)數(shù)值精度 FP32 的模型，一部分算子的操作時(shí)，數(shù)值精度為 FP16，其余算子的操作精度是 FP32，而具體哪些算子用 FP16，哪些用 FP32，不需要用戶關(guān)心，amp 自動(dòng)給它們都安排好了。這樣在不改變模型、不降低模型訓(xùn)練精度的前提下，可以縮短訓(xùn)練時(shí)間，降低存儲(chǔ)需求，因而能支持更多的 batch size、更大模型和尺寸更大的輸入進(jìn)行訓(xùn)練。PyTorch 從 1.6 以后（在此之前 OpenMMLab 已經(jīng)支持混合精度訓(xùn)練，即 Fp16OptimizerHook），開始原生支持 amp，即torch.cuda.amp module。2020 ECCV，英偉達(dá)官方做了一個(gè)?tutorial?推廣 amp。從官方各種文檔網(wǎng)頁 claim 的結(jié)果來看，amp 在分類、檢測、圖像生成、3D CNNs、LSTM，以及 NLP 中機(jī)器翻譯、語義識(shí)別等應(yīng)用中，都在沒有降低模型訓(xùn)練精度都前提下，加速了模型的訓(xùn)練速度。
本文是對(duì)torch.cuda.amp工作機(jī)制，和 module 中接口使用方法介紹，以及在算法角度上對(duì) amp 不掉點(diǎn)原因進(jìn)行分析，最后補(bǔ)充一點(diǎn)對(duì) amp 存儲(chǔ)消耗的解釋。
1. 混合精度訓(xùn)練機(jī)制
torch.cuda.amp?給用戶提供了較為方便的混合精度訓(xùn)練機(jī)制，“方便”體現(xiàn)在兩個(gè)方面：
用戶不需要手動(dòng)對(duì)模型參數(shù) dtype 轉(zhuǎn)換，amp 會(huì)自動(dòng)為算子選擇合適的數(shù)值精度
對(duì)于反向傳播的時(shí)候，FP16 的梯度數(shù)值溢出的問題，amp 提供了梯度 scaling 操作，而且在優(yōu)化器更新參數(shù)前，會(huì)自動(dòng)對(duì)梯度 unscaling，所以，對(duì)用于模型優(yōu)化的超參數(shù)不會(huì)有任何影響
以上兩點(diǎn)，分別是通過使用amp.autocast和amp.GradScaler來實(shí)現(xiàn)的。
autocast可以作為 Python 上下文管理器和裝飾器來使用，用來指定腳本中某個(gè)區(qū)域、或者某些函數(shù)，按照自動(dòng)混合精度來運(yùn)行。混合精度在操作的時(shí)候，是先將 FP32 的模型的參數(shù)拷貝一份，拷貝的參數(shù)轉(zhuǎn)換成 FP16，而 amp 規(guī)定了的 FP16 的算子（例如卷積、全連接），對(duì) FP16 的數(shù)值進(jìn)行操作；FP32 的算子（例如涉及 reduction 的算子，BatchNormalize，softmax...），輸入和輸出是 FP16，計(jì)算的精度是 FP32。在反向傳播時(shí)，依然是混合精度計(jì)算，得到數(shù)值精度為 FP16 的梯度。最后，由于 GPU 中的 Tensor Core 天然支持 FP16 乘積的結(jié)果與 FP32 的累加（Tensor Core math），優(yōu)化器的操作是利用 FP16 的梯度對(duì) FP32 的參數(shù)進(jìn)行更新。
對(duì)于 FP16 不可避免的問題就是：表示的范圍較窄，如下圖所示，大量非 0 梯度會(huì)遇到溢出問題。解決辦法是：對(duì)梯度乘一個(gè)?
?的系數(shù)，稱為 scale factor，把梯度 shift 到 FP16 的表示范圍。
GradScaler的工作就是在反向傳播前給 loss 乘一個(gè) scale factor，所以之后反向傳播得到的梯度都乘了相同的 scale factor。并且為了不影響學(xué)習(xí)率，在梯度更新前將梯度unscale?？偨Y(jié)amp的基本訓(xùn)練流程：
維護(hù)一個(gè) FP32 數(shù)值精度模型的副本
在每個(gè)iteration
拷貝并且轉(zhuǎn)換成 FP16 模型
前向傳播（FP16 的模型參數(shù)）
loss 乘 scale factor?
s反向傳播（FP16 的模型參數(shù)和參數(shù)梯度）
參數(shù)梯度乘?
1/s利用 FP16 的梯度更新 FP32 的模型參數(shù)
但是，這里會(huì)有一個(gè)問題，scale factor 應(yīng)該如何選取？選一個(gè)常量顯然是不合適的，因?yàn)?loss 和梯度的數(shù)值在變，scale factor 需要跟隨 loss 動(dòng)態(tài)變化。健康的 loss 是振蕩中下降，因此GradScaler設(shè)計(jì)的 scale factor 每隔?
?個(gè) iteration 乘一個(gè)大于 1 的系數(shù)，再 scale loss；并且每次更新前檢查溢出問題（檢查梯度中有沒有inf和nan），如果有，scale factor 乘一個(gè)小于 1 的系數(shù)并跳過該 iteration 的參數(shù)更新環(huán)節(jié)，如果沒有，就正常更新參數(shù)。動(dòng)態(tài)更新 scale factor 是 amp 實(shí)際操作中的流程?？偨Y(jié) amp 動(dòng)態(tài) scale factor 的訓(xùn)練流程：
維護(hù)一個(gè) FP32 數(shù)值精度模型的副本
初始化?
s在每個(gè) iteration + a 拷貝并且轉(zhuǎn)換成FP16模型 + b 前向傳播（FP16 的模型參數(shù)） + c loss 乘 scale factor?
s?+ d 反向傳播（FP16 的模型參數(shù)和參數(shù)梯度） + e 檢查有沒有inf或者nan的參數(shù)梯度 + 如果有：降低?s，回到步驟a + f 參數(shù)梯度乘?1/s?+ g 利用 FP16 的梯度更新 FP32 的模型參數(shù)
2. amp模塊的API
用戶使用混合精度訓(xùn)練基本操作：
# amp依賴Tensor core架構(gòu)，所以model參數(shù)必須是cuda tensor類型
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# GradScaler對(duì)象用來自動(dòng)做梯度縮放
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 在autocast enable 區(qū)域運(yùn)行forward
with autocast():
# model做一個(gè)FP16的副本，forward
output = model(input)
loss = loss_fn(output, target)
# 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16)
scaler.scale(loss).backward()
# scaler 更新參數(shù)，會(huì)先自動(dòng)unscale梯度
# 如果有nan或inf，自動(dòng)跳過
scaler.step(optimizer)
# scaler factor更新
scaler.update()2.1 autocast類
autocast(enable=True)`` 可以作為上下文管理器和裝飾器來使用，給算子自動(dòng)安排按照 FP16 或者 FP32 的數(shù)值精度來操作。
2.1.1 autocast算子
PyTorch中，只有 CUDA 算子有資格被 autocast，而且只有 “out-of-place” 才可以被 autocast，例如：a.addmm(b, c)是可以被 autocast，但是a.addmm_(b, c)和a.addmm(b, c, out=d)不可以 autocast。amp autocast 成 FP16 的算子有：
autocast 成 FP32 的算子：
剩下沒有列出的算子，像dot,add,cat...都是按數(shù)據(jù)中較大的數(shù)值精度，進(jìn)行操作，即有 FP32 參與計(jì)算，就按 FP32，全是 FP16 參與計(jì)算，就是 FP16。
2.1.2 MisMatch error
作為上下文管理器使用時(shí)，混合精度計(jì)算 enable 區(qū)域得到的 FP16 數(shù)值精度的變量在 enable 區(qū)域外需要顯式的轉(zhuǎn)成 FP32：
# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")
with autocast():
# torch.mm is on autocast's list of ops that should run in float16.
e_float16 = torch.mm(a_float32, b_float32)
# Also handles mixed input types
f_float16 = torch.mm(d_float32, e_float16)
# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())2.1.3 autocast 嵌套使用
# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")
with autocast():
e_float16 = torch.mm(a_float32, b_float32)
with autocast(enabled=False):
f_float32 = torch.mm(c_float32, e_float16.float())
g_float16 = torch.mm(d_float32, f_float32)2.1.4 autocast 作為裝飾器
這種情況一般用于 data parallel 的模型的，autocast 設(shè)計(jì)為 “thread local” 的，所以只在 main thread 上設(shè) autocast 區(qū)域是不 work 的:
model = MyModel()
dp_model = nn.DataParallel(model)
with autocast(): # dp_model's internal threads won't autocast.
#The main thread's autocast state has no effect.
output = dp_model(input) # loss_fn still autocasts, but it's too late...
loss = loss_fn(output) 正確姿勢是對(duì) forward 裝飾：
MyModel(nn.Module):
...
@autocast()
def forward(self, input):
...另一個(gè)正確姿勢是在 forward 的里面設(shè) autocast 區(qū)域：
MyModel(nn.Module):
...
def forward(self, input):
with autocast():
...forward 函數(shù)處理之后，在 main thread 里 autocast
model = MyModel()
dp_model = nn.DataParallel(model)
with autocast():
output = dp_model(input)
loss = loss_fn(output)2.1.5 autocast 自定義函數(shù)
對(duì)于用戶自定義的 autograd 函數(shù)，需要用amp.custom_fwd裝飾 forward 函數(shù)，amp.custom_bwd裝飾 backward 函數(shù)：
class MyMM(torch.autograd.Function):
@staticmethod
@custom_fwd
def forward(ctx, a, b):
ctx.save_for_backward(a, b)
return a.mm(b)
@staticmethod
@custom_bwd
def backward(ctx, grad):
a, b = ctx.saved_tensors
return grad.mm(b.t()), a.t().mm(grad)調(diào)用時(shí)再 autocast
mymm = MyMM.apply
with autocast():
output = mymm(input1, input2)2.1.6 源碼分析
autocast主要實(shí)現(xiàn)接口有：
A.?__enter__
def __enter__(self):
self.prev = torch.is_autocast_enabled()
torch.set_autocast_enabled(self._enabled)
torch.autocast_increment_nesting()B.?__exit__
def __exit__(self, *args):
if torch.autocast_decrement_nesting() == 0:
torch.clear_autocast_cache()
torch.set_autocast_enabled(self.prev)
return FalseC.?__call__
def __call__(self, func):
@functools.wraps(func)
def decorate_autocast(*args, **kwargs):
with self:
return func(*args, **kwargs)
return decorate_autocast其中torch.*autocast*函數(shù)是在 pytorch/aten/src/ATen/autocast_mode.cpp 里實(shí)現(xiàn)。PyTorch ATen 是 A TENsor library for C++11，ATen 部分有大量的代碼是來聲明和定義 Tensor 運(yùn)算相關(guān)的邏輯的。autocast_mode.cpp?實(shí)現(xiàn)策略是 “ cache fp16 casts of fp32 model weights”。
2.2 GradScaler 類
torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)用于動(dòng)態(tài) scale 梯度
+.?init_scale: scale factor 的初始值 +.?growth_factor: 每次 scale factor 的增長系數(shù) +.?backoff_factor: scale factor 下降系數(shù) +.?growth_interval: 每隔多個(gè) interval 增長 scale factor +.?enabled: 是否做 scale
2.2.1?scale(output)方法
對(duì)outputs乘 scale factor，并返回，如果enabled=False就原樣返回。
2.2.3?step(optimizer, *args, **kwargs)方法
step?方法在做兩件事情：
對(duì)梯度 unscale，如果之前沒有手動(dòng)調(diào)用
unscale方法的話檢查梯度溢出，如果沒有
nan/inf，就執(zhí)行 optimizer 的 step，如果有就跳過
注意：GradScaler的step不支持傳 closure。
2.2.4?update(new_scale=None)方法
update方法在每個(gè) iteration 結(jié)束前都需要調(diào)用，如果參數(shù)更新跳過，會(huì)給 scale factor 乘backoff_factor，或者到了該增長的 iteration，就給 scale factor 乘growth_factor。也可以用new_scale直接更新 scale factor。
2.3 舉例
2.3.1 Gradient clipping
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast():
output = model(input)
loss = loss_fn(output, target)
scaler.scale(loss).backward()
# unscale 梯度，可以不影響clip的threshold
scaler.unscale_(optimizer)
# clip梯度
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
# unscale_（）已經(jīng)被顯式調(diào)用了，scaler正常執(zhí)行step更新參數(shù)，有nan/inf也會(huì)跳過
scaler.step(optimizer)
scaler.update()2.3.2 Gradient accumulation
scaler = GradScaler()
for epoch in epochs:
for i, (input, target) in enumerate(data):
with autocast():
output = model(input)
loss = loss_fn(output, target)
# loss 根據(jù) 累加的次數(shù)歸一一下
loss = loss / iters_to_accumulate
# scale 歸一的loss 并backward
scaler.scale(loss).backward()
if (i + 1) % iters_to_accumulate == 0:
# may unscale_ here if desired
# (e.g., to allow clipping unscaled gradients)
# step() and update() proceed as usual.
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()2.3.3. Gradient penalty
scaler = GradScaler()
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 防止溢出，在不是autocast 區(qū)域，先用scaled loss 得到 scaled 梯度
scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
inputs=model.parameters(),
create_graph=True)
# 梯度unscale
inv_scale = 1./scaler.get_scale()
grad_params = [p * inv_scale for p in scaled_grad_params]
# 在autocast 區(qū)域，loss 加上梯度懲罰項(xiàng)
with autocast():
grad_norm = 0
for grad in grad_params:
grad_norm += grad.pow(2).sum()
grad_norm = grad_norm.sqrt()
loss = loss + grad_norm
scaler.scale(loss).backward()
# may unscale_ here if desired
# (e.g., to allow clipping unscaled gradients)
# step() and update() proceed as usual.
scaler.step(optimizer)
scaler.update()2.3.4. Multiple models
scaler 一個(gè)就夠，但 scale(loss) 和 step(optimizer) 要分別執(zhí)行
scaler = torch.cuda.amp.GradScaler()
for epoch in epochs:
for input, target in data:
optimizer0.zero_grad()
optimizer1.zero_grad()
with autocast():
output0 = model0(input)
output1 = model1(input)
loss0 = loss_fn(2 * output0 + 3 * output1, target)
loss1 = loss_fn(3 * output0 - 5 * output1, target)
# (retain_graph here is unrelated to amp, it's present because in this
# example, both backward() calls share some sections of graph.)
scaler.scale(loss0).backward(retain_graph=True)
scaler.scale(loss1).backward()
# You can choose which optimizers receive explicit unscaling, if you
# want to inspect or modify the gradients of the params they own.
scaler.unscale_(optimizer0)
scaler.step(optimizer0)
scaler.step(optimizer1)
scaler.update()2.3.5. Multiple GPUs
torch DDP 和 torch DP model 的處理方式一樣
Q1. amp 是如何做到 FP16 和 FP32 混合使用，“還不掉點(diǎn)”
模型量化、模型壓縮的算法挺多的，但都做不 amp 這樣，對(duì)多數(shù)模型訓(xùn)練不掉點(diǎn)（但是實(shí)操中，聽有經(jīng)驗(yàn)的大神介紹，完全不到點(diǎn)還是有點(diǎn)難度的）。amp 能做成這樣，因?yàn)樗鼘?duì)模型沒有壓縮和量化，維護(hù)的還是一個(gè) 32 位的模型。只是用 16 位去表示原來 32 位的梯度：通常模型訓(xùn)練依賴 FP32 的精度，因?yàn)樘荻葧?huì)有一部分 FP16 表示不了，而 scale factor 把梯度 shift 到 FP16 能表示范圍，使得梯度方面精度的損失較小，可能 forward 時(shí)候的直接的精度壓縮是訓(xùn)練最大的損失。
Q2. 沒有 Tensor Core 架構(gòu)能否使用 amp
沒有 Tensor Core 架構(gòu)的 GPU 試用 amp，速度反而下降，但顯存會(huì)明顯減少。作者在 Turing 架構(gòu)的 GTX 1660 上試用 amp，運(yùn)算時(shí)間增加了一倍，但顯存不到原來的一半。
Q3. 為什么 amp 中有兩份參數(shù)，存儲(chǔ)消耗反而更小
相比與模型參數(shù)，對(duì)中間層結(jié)果的存儲(chǔ)更是 deep learning 的 bottleneck。當(dāng)對(duì)中間結(jié)果的存儲(chǔ)砍半，整個(gè)存儲(chǔ)消耗就基本上原來的一半。
推薦閱讀
谷歌提出Meta Pseudo Labels，刷新ImageNet上的SOTA！
漲點(diǎn)神器FixRes：兩次超越ImageNet數(shù)據(jù)集上的SOTA
SWA：讓你的目標(biāo)檢測模型無痛漲點(diǎn)1% AP
CondInst：性能和速度均超越Mask RCNN的實(shí)例分割模型
mmdetection最小復(fù)刻版(十一)：概率Anchor分配機(jī)制PAA深入分析
MMDetection新版本V2.7發(fā)布，支持DETR，還有YOLOV4在路上!
無需tricks，知識(shí)蒸餾提升ResNet50在ImageNet上準(zhǔn)確度至80%+
不妨試試MoCo，來替換ImageNet上pretrain模型！
mmdetection最小復(fù)刻版(七)：anchor-base和anchor-free差異分析
mmdetection最小復(fù)刻版(四)：獨(dú)家yolo轉(zhuǎn)化內(nèi)幕
機(jī)器學(xué)習(xí)算法工程師
? ??? ? ? ? ? ? ? ? ? ? ? ??????? ??一個(gè)用心的公眾號(hào)
?
