簡介

自動混合精度訓(xùn)練（auto Mixed Precision，amp）是深度學(xué)習(xí)比較流行的一個訓(xùn)練技巧，它可以大幅度降低訓(xùn)練的成本并提高訓(xùn)練的速度，因此在競賽中受到了較多的關(guān)注。此前，比較流行的混合精度訓(xùn)練工具是由NVIDIA開發(fā)的A PyTorch Extension(Apex)，它能夠以非常簡單的API支持自動混合精度訓(xùn)練，不過，PyTorch從1.6版本開始已經(jīng)內(nèi)置了amp模塊，本文簡單介紹其使用。

自動混合精度（AMP）

首先來聊聊自動混合精度的由來。下圖是常見的浮點數(shù)表示形式，它表示單精度浮點數(shù)，在編程語言中的體現(xiàn)是float型，顯然從圖中不難看出它需要4個byte也就是32bit來進(jìn)行存儲。深度學(xué)習(xí)的模型數(shù)據(jù)均采用float32進(jìn)行表示，這就帶來了兩個問題：模型size大，對顯存要求高；32位計算慢，導(dǎo)致模型訓(xùn)練和推理速度慢。

那么半精度是什么呢，顧名思義，它只用16位即2byte來進(jìn)行表示，較小的存儲占用以及較快的運算速度可以緩解上面32位浮點數(shù)的兩個主要問題，因此半精度會帶來下面的一些優(yōu)勢：

1. 顯存占用更少，模型只有32位的一半存儲占用，這也可以使用更大的batch size以適應(yīng)一些對大批尺寸有需求的結(jié)構(gòu)，如Batch Normalization；
2. 計算速度快，float16的計算吞吐量可以達(dá)到float32的2-8倍左右，且隨著NVIDIA張量核心的普及，使用半精度計算已經(jīng)比較成熟，它會是未來深度學(xué)習(xí)計算的一個重要趨勢。

那么，半精度有沒有什么問題呢？其實也是有著很致命的問題的，主要是移除錯誤和舍入誤差兩個方面，具體可以參考這篇文章，作者解析的很好，我這里就簡單復(fù)述一下。

溢出錯誤

FP16的數(shù)值表示范圍比FP32的表示范圍小很多，因此在計算過程中很容易出現(xiàn)上溢出（overflow）和下溢出（underflow）問題，溢出后會出現(xiàn)梯度nan問題，導(dǎo)致模型無法正確更新，嚴(yán)重影響網(wǎng)絡(luò)的收斂。而且，深度模型訓(xùn)練，由于激活函數(shù)的梯度往往比權(quán)重的梯度要小，更容易出現(xiàn)的是下溢出問題。

舍入誤差

舍入誤差（Rounding Error）指的是當(dāng)梯度過小，小于當(dāng)前區(qū)間內(nèi)的最小間隔時，該次梯度更新可能會失敗。上面說的知乎文章的作者用來一張很形象的圖進(jìn)行解釋，具體如下，意思是說在2?32^{-3}2?3到2?22^{-2}2?2之間，2?32^{-3}2?3每次變大都會至少加上2?132^{-13}2?13，顯然，梯度還在這個間隔內(nèi)，因此更新是失敗的。

那么這兩個問題是如何解決的呢，思路來自于NVIDIA和百度合作的論文，我這里簡述一下方法：混合精度訓(xùn)練和損失縮放。前者的思路是在內(nèi)存中使用FP16做儲存和乘法運算以加速計算，用FP32做累加運算以避免舍入誤差，這樣就緩解了舍入誤差的問題；后者則是針對梯度值太小從而下溢出的問題，它的思想是：反向傳播前，將損失變化手動增大2k2^k2k倍，因此反向傳播時得到的中間變量（激活函數(shù)梯度）則不會溢出；反向傳播后，將權(quán)重梯度縮小2k2^k2k倍，恢復(fù)正常值。

研究人員通過引入FP32進(jìn)行混合精度訓(xùn)練以及通過損失縮放來解決FP16的不足，從而實現(xiàn)了一套混合精度訓(xùn)練的范式，NVIDIA以此為基礎(chǔ)設(shè)計了Apex包，不過Apex的使用本文就不涉及了，下一節(jié)主要關(guān)注如何使用torch.cuda.amp實現(xiàn)自動混合精度訓(xùn)練，不過這里還需要補充的一點就是目前混合精度訓(xùn)練支持的N卡只有包含Tensor Core的卡，如2080Ti、Titan、Tesla等。

PyTorch自動混合精度

PyTorch對混合精度的支持始于1.6版本，位于torch.cuda.amp模塊下，主要是torch.cuda.amp.autocast和torch.cuda.amp.GradScale兩個模塊，autocast針對選定的代碼塊自動選取適合的計算精度，以便在保持模型準(zhǔn)確率的情況下最大化改善訓(xùn)練效率；GradScaler通過梯度縮放，以最大程度避免使用FP16進(jìn)行運算時的梯度下溢。官方給的使用這兩個模塊進(jìn)行自動精度訓(xùn)練的示例代碼鏈接給出，我對其示例解析如下，這就是一般的訓(xùn)練框架。

# 以默認(rèn)精度創(chuàng)建模型和優(yōu)化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)# 創(chuàng)建梯度縮放器
scaler = GradScaler()for epoch in epochs:for input, target in data:optimizer.zero_grad()# 通過自動類型轉(zhuǎn)換進(jìn)行前向傳播with autocast():output = model(input)loss = loss_fn(output, target)# 縮放大損失，反向傳播不建議放到autocast下，它默認(rèn)和前向采用相同的計算精度scaler.scale(loss).backward()# 先反縮放梯度，若反縮后梯度不是inf或者nan，則用于權(quán)重更新scaler.step(optimizer)# 更新縮放器scaler.update()

下面我以簡單的MNIST任務(wù)做測試，使用的顯卡為RTX 3090，代碼如下。該代碼段中只包含核心的訓(xùn)練模塊，模型的定義和數(shù)據(jù)集的加載熟悉PyTorch的應(yīng)該不難自行補充。

model = Model()
model = model.cuda()
loss_fn = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())n_epochs = 30
start = time.time()
for epoch in range(n_epochs):total_loss, correct, total = 0.0, 0, 0model.train()for step, data in enumerate(data_loader_train):x_train, y_train = datax_train, y_train = x_train.cuda(), y_train.cuda()outputs = model(x_train)_, pred = torch.max(outputs, 1)loss = loss_fn(outputs, y_train)optimizer.zero_grad()loss.backward()optimizer.step()total_loss += loss.item()total += len(y_train)correct += torch.sum(pred == y_train).item()print("epoch {} loss {} acc {}".format(epoch, total_loss, correct / total))

我這里采用的是一個很小的模型，又是一個很簡單的任務(wù)，因此模型都是很快收斂，因此精度上沒有什么明顯的區(qū)別，不過如果是訓(xùn)練大型模型的話，有人已經(jīng)用實驗證明，內(nèi)置amp和apex庫都會有精度下降，不過amp效果更好一些，下降較少。上面的loss變化圖也是非常類似的。

再來看存儲方面，顯存縮減在這個任務(wù)中的表現(xiàn)不是特別明顯，因為這個任務(wù)的參數(shù)量不多，前后向過程中的FP16存儲節(jié)省不明顯，而因為引入了一些拷貝之類的，反而使得顯存略有上升，實際的任務(wù)中，這種開銷肯定遠(yuǎn)小于FP32的開銷的。

最后，不妨看一下使用混合精度最關(guān)心的速度問題，實際上混合精度確實會帶來一些速度上的優(yōu)勢，一些官方的大模型如BERT等訓(xùn)練速度提高了2-3倍，這對于工業(yè)界的需求來說，啟發(fā)還是比較多的。

總結(jié)

混合精度計算是未來深度學(xué)習(xí)發(fā)展的重要方向，很受工業(yè)界的關(guān)注，PyTorch從1.6版本開始默認(rèn)支持amp，雖然現(xiàn)在還不是特別完善，但以后一定會越來越好，因此熟悉自動混合精度的用法還是有必要的。

超強干貨來襲 云風(fēng)專訪：近40年碼齡，通宵達(dá)旦的技術(shù)人生

總結(jié)

以上是生活随笔為你收集整理的PyTorch-混合精度训练的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。