本文將通過一個(gè)實(shí)戰(zhàn)案例來展示如何借助于PyTorch自動(dòng)混合精度庫對ResNet50模型進(jìn)行優(yōu)化，然后借助少許幾行代碼即可獲得超過兩倍速度的模型訓(xùn)練效率。

簡介

你是否曾希望你的深度學(xué)習(xí)模型運(yùn)行得更快？

一方面，GPU很昂貴。另一方面，數(shù)據(jù)集龐大，訓(xùn)練過程似乎永無止境；你可能有一百萬個(gè)實(shí)驗(yàn)要進(jìn)行，還有一個(gè)截止日期。所有這些需求都是期待特定形式的訓(xùn)練加速的極好理由。

但是，我們該選哪一種模型呢？

??PyTorch???、??HuggingFace???和??Nvidia??已經(jīng)為模型訓(xùn)練的性能調(diào)優(yōu)提供了很好的參考，包括異步數(shù)據(jù)加載、緩沖區(qū)檢查點(diǎn)、分布式數(shù)據(jù)并行化和自動(dòng)混合精度等等。

在這篇文章中，我將專注介紹自動(dòng)混合精度技術(shù)。首先，我將簡要介紹Nvidia的張量核設(shè)計(jì)；然后，我們一起探討發(fā)表在ICLR 2018上的開創(chuàng)性工作——“混合精度訓(xùn)練”相關(guān)論文；最后，我將介紹一個(gè)在FashionMNIST數(shù)據(jù)集上訓(xùn)練ResNet50模型的簡單示例。通過這個(gè)示例，我們來展示如何在加載雙倍批量數(shù)據(jù)的同時(shí)將訓(xùn)練速度提高兩倍，而這一結(jié)果卻只需要額外編寫三行代碼。

硬件基礎(chǔ)——Nvidia張量核

首先，讓我們回顧一下GPU設(shè)計(jì)的一些基本原理。英偉達(dá)GPU最受歡迎的商業(yè)產(chǎn)品之一是Volta系列，例如基于GV100 GPU設(shè)計(jì)的V100 GPU。因此，我們將圍繞下面的GV100架構(gòu)進(jìn)行討論。

對于GV100架構(gòu)來說，流式多處理器（SM）是計(jì)算的核心設(shè)計(jì)。每個(gè)GPU包含6個(gè)GPU處理集群（GPC）和S84 SM（或V100的80 SM）。其整體設(shè)計(jì)如下圖所示：

Volta GV100 GPU設(shè)計(jì)（每個(gè)GPU包含6個(gè)GPC，每個(gè)GPC包含14個(gè)SM；圖像來源：https://arxiv.org/pdf/1803.04014）

對于每個(gè)SM，它包含兩種類型的核心：CUDA核心和張量核。CUDA核心是Nvidia于2006年推出的原始設(shè)計(jì)，是CUDA平臺(tái)的重要組成部分。CUDA核心可分為三種類型：FP64核心/單元、FP32核心/單元和Int32核心/單元。每個(gè)GV100 SM包含32個(gè)FP64核心、64個(gè)FP32核心和64個(gè)Int32核心。Volta/Turing（2017）系列GPU中引入了張量核，以便與之前的Pascal（2016）系列分離。GV100上的每個(gè)SM包含8個(gè)張量核。??鏈接??處給出了V100 GPU的完整詳細(xì)信息列表。下面詳細(xì)介紹SM設(shè)計(jì)。

一個(gè)英偉達(dá)Tesla V100處理組就包含640個(gè)張量核（圖像來源：https://arxiv.org/pdf/1903.03640）

為什么選擇張量核？Nvidia張量核專門用于執(zhí)行通用矩陣乘法（GEMM）和半精度矩陣乘法和累加（HMMA）操作。簡而言之，GEMM以A*B+C的格式執(zhí)行矩陣運(yùn)算，HMMA將運(yùn)算轉(zhuǎn)換為半精度格式。有關(guān)這方面的更詳細(xì)的討論可以在??鏈接???處找到。由于深度學(xué)習(xí)涉及MMA；所以，張量（Tensor）核心在當(dāng)今的模型訓(xùn)練和加速計(jì)算中至關(guān)重要。

GEMM操作示例（對于HMMA，A和B通常轉(zhuǎn)換為FP16，而C和D可以是FP16或FP32），圖像來源：https://arxiv.org/pdf/1811.08309

當(dāng)然，當(dāng)切換到混合精度訓(xùn)練時(shí)，請務(wù)必檢查你使用的GPU的規(guī)格。只有最新的GPU系列支持張量核，混合精度訓(xùn)練只能在這些機(jī)器上使用。

數(shù)據(jù)格式基礎(chǔ)——單精度（FP32）與半精度（FP16）

現(xiàn)在，讓我們仔細(xì)看看FP32和FP16格式。FP32和FP16是IEEE格式，使用32位二進(jìn)制存儲(chǔ)和16位二進(jìn)制存儲(chǔ)表示浮點(diǎn)數(shù)。這兩種格式都包括三個(gè)部分：a）符號(hào)位；b）指數(shù)位；c）尾數(shù)位。FP32和FP16分配給指數(shù)和尾數(shù)的比特?cái)?shù)不同，這導(dǎo)致了不同的值范圍和精度。

FP16（IEEE標(biāo)準(zhǔn)）、BF16（Google Brain標(biāo)準(zhǔn)）、FP32（IEEE標(biāo)準(zhǔn)）和TF32（Nvidia標(biāo)準(zhǔn)）之間的區(qū)別（圖像來源：https://en.wikipedia.org/wiki/Bfloat16_floating-point_format）

如何將FP16和FP32轉(zhuǎn)換為真實(shí)的值呢？根據(jù)IEEE-754標(biāo)準(zhǔn)，F(xiàn)P32的十進(jìn)制值=（-1）^（符號(hào)）×2^（十進(jìn)制指數(shù)-127）×（隱式前導(dǎo)1+十進(jìn)制尾數(shù)），其中127是有偏差的指數(shù)值。對于FP16，公式變?yōu)椋?1）^（符號(hào)）×2^（十進(jìn)制指數(shù)-15）×（隱式前導(dǎo)1+十進(jìn)制尾數(shù)），其中15是相應(yīng)的有偏指數(shù)值。你可以在??鏈接??處查看有偏指數(shù)值的更多詳細(xì)信息。

從這個(gè)意義上講，F(xiàn)P32的取值范圍約為[-212S，21283;]~[-1.7*1e38，1.7*1e38]，F(xiàn)P16的取值范圍大約為[-2?，2'8309]=[-32768，32768]。請注意，F(xiàn)P32的十進(jìn)制指數(shù)在0到255之間，我們排除了最大值0xFF，因?yàn)樗硎綨AN。這就解釋了為什么最大的十進(jìn)制指數(shù)是254–127=127。當(dāng)然，類似的規(guī)則也適用于FP16。

對于精度方面，請注意指數(shù)和尾數(shù)都有助于精度限制（也稱為非規(guī)范化，請參閱??鏈接???處的詳細(xì)討論）。因此，F(xiàn)P32可以表示高達(dá)2^（-23）*2^（-126）=2^（-149）的精度，F(xiàn)P16可以表示高至2^（10）*2^。

FP32和FP16表示之間的差異帶來了混合精度訓(xùn)練的關(guān)鍵問題，因?yàn)樯疃葘W(xué)習(xí)模型的不同層/操作對值范圍和精度或者不敏感或者敏感，所以需要單獨(dú)解決。

混合精度訓(xùn)練

前面，我們已經(jīng)學(xué)習(xí)了MMA的硬件基礎(chǔ)知識(shí)、張量核的概念以及FP32和FP16之間的關(guān)鍵區(qū)別。接下來，我們可以進(jìn)一步討論混合精度訓(xùn)練的細(xì)節(jié)。

混合精度訓(xùn)練的想法最早是在2018年ICLR論文??《混合精度訓(xùn)練》（Mixed Precision Training）???中提出的。該論文在訓(xùn)練過程中將深度學(xué)習(xí)模型轉(zhuǎn)換為半精度浮點(diǎn)，而不會(huì)損失模型精度或修改超參數(shù)。如上所述，由于FP32和FP16之間的關(guān)鍵區(qū)別在于值范圍和精度，該論文詳細(xì)討論了FP16為什么會(huì)導(dǎo)致梯度消失，以及如何通過損失縮放來解決這個(gè)問題。此外，該論文還提出了使用FP32主權(quán)重拷貝和使用FP32進(jìn)行歸約和向量點(diǎn)積累加等特定操作的技巧。

損失縮放（Loss scaling）。本文給出了一個(gè)使用FP32精度訓(xùn)練Multibox SSD探測器網(wǎng)絡(luò)的示例，如下所示。如果不進(jìn)行任何縮放，F(xiàn)P16梯度的指數(shù)范圍≥2^（-24），以下所有值都將變?yōu)榱?，這與FP32相比是不夠的。然而，通過實(shí)驗(yàn)，將梯度簡單地縮放23=8倍，可以使半精度訓(xùn)練精度恢復(fù)到與FP32相匹配的水平。從這個(gè)意義上講，作者認(rèn)為[2^（-27），2^（-24）]之間的額外百分之幾的梯度在訓(xùn)練過程中仍然很重要，而低于2^（-27）的值并不重要。

在Multibox SSD訓(xùn)練示例中使用FP32精度的梯度值范圍。請注意，[2^（-27），2^（-24）]之間的值超出了FP16非規(guī)范化范圍，僅占總梯度的百分之幾，但在整體訓(xùn)練中仍然很重要。圖像來源：https://arxiv.org/pdf/1710.03740

解決這種規(guī)模差異的方法是借助損失縮放的辦法。根據(jù)鏈?zhǔn)椒▌t，縮放損失將確保相同的量將縮放所有梯度。但是請注意，在最終權(quán)重更新之前，需要取消縮放梯度。

自動(dòng)混合精度訓(xùn)練

Nvidia公司首先開發(fā)了名為APEX的PyTorch擴(kuò)展自動(dòng)混合精度訓(xùn)練，然后被PyTorch、TensorFlow、MXNet等主流框架廣泛采用。請參閱??鏈接??處的Nvidia文檔。為了簡單起見，我們只介紹PyTorch框架中的自動(dòng)混合精度庫：https://pytorch.org/docs/stable/amp.html.

amp庫可以自動(dòng)處理大多數(shù)混合精度訓(xùn)練技術(shù)，如FP32主權(quán)重復(fù)制。開發(fā)人員只需要進(jìn)行操作數(shù)自動(dòng)轉(zhuǎn)換和梯度/損失縮放。

操作數(shù)自動(dòng)轉(zhuǎn)換：盡管我們提到張量核可以大大提高GEMM操作的性能，但某些操作不適合半精度表示。

amp庫給出了一個(gè)符合半精度條件的??CUDA操作列表??。amp.autocast完全涵蓋了大多數(shù)矩陣乘法、卷積和線性激活運(yùn)算；但是，對于歸約/求和、softmax和損失計(jì)算等，這些計(jì)算仍然在FP32中執(zhí)行，因?yàn)樗鼈儗?shù)據(jù)范圍和精度更敏感。

梯度/損失縮放：amp庫提供了??自動(dòng)梯度縮放技術(shù)???；因此，用戶在訓(xùn)練過程中不必手動(dòng)調(diào)整縮放。??鏈接??處可以找到縮放因子的更詳細(xì)的算法。

一旦縮放了梯度，就需要在進(jìn)行梯度剪裁和正則化之前將其縮小。更多細(xì)節(jié)可以在??鏈接??處找到。

FashionMNIST訓(xùn)練示例

torch.amp庫相對易于使用，只需要三行代碼即可將訓(xùn)練速度提高2倍。

首先，我們從一個(gè)非常簡單的任務(wù)開始，使用FP32在FashionMNIST數(shù)據(jù)集（MIT許可證）上訓(xùn)練ResNet50模型；我們可以看到10個(gè)世代的訓(xùn)練時(shí)間為333秒：

ResNet50模型在數(shù)據(jù)集FashionMNIST上的訓(xùn)練

小于2**（-24）的梯度與總梯度之比。我們可以看到，F(xiàn)P16將使總梯度的近1/4變?yōu)榱?/p>

評(píng)估結(jié)果

現(xiàn)在，我們使用amp庫。amp庫只需要三行額外的代碼進(jìn)行混合精度訓(xùn)練。我們可以看到訓(xùn)練在141秒內(nèi)完成，比FP32訓(xùn)練快2.36倍，同時(shí)達(dá)到了相同的精確度、召回率和F1分?jǐn)?shù)。

scaler = torch.cuda.amp.GradScaler()
#開始訓(xùn)練的代碼
# ...
with torch.autocast(device_type="cuda"):
  #訓(xùn)練代碼

#封裝損失與優(yōu)化器
scaler.scale(loss).backward()
scaler.step(optimizer)

scaler.update()

使用amp庫進(jìn)行訓(xùn)練

訓(xùn)練期間的縮放因子（縮放因子僅在第一步發(fā)生變化，保持不變。）

最終結(jié)果與FP32訓(xùn)練結(jié)果比較

上面代碼的Github鏈接：??https://github.com/adoskk/MachineLearningBasics/blob/main/mixed_precision_training/mixed_precision_training.ipynb。???

總結(jié)

混合精度訓(xùn)練是加速深度學(xué)習(xí)模型訓(xùn)練的一種非常有價(jià)值的技術(shù)。它不僅加快了浮點(diǎn)運(yùn)算的速度，還節(jié)省了GPU內(nèi)存，因?yàn)橛?xùn)練批次可以轉(zhuǎn)換為FP16，從而節(jié)省了一半的GPU內(nèi)存。另外，借助于PyTorch框架中的amp庫，額外的代碼可以減少到僅僅三行，因?yàn)闄?quán)重復(fù)制、損失縮放、操作類型轉(zhuǎn)換等計(jì)算都是由該庫內(nèi)部處理的。

需要注意的是，如果模型權(quán)重大小遠(yuǎn)大于數(shù)據(jù)批次的話，混合精度訓(xùn)練并不能真正解決GPU內(nèi)存問題。首先，只有模型的某些層被轉(zhuǎn)換成FP16，而其余層仍在FP32中計(jì)算；其次，權(quán)重更新需要FP32復(fù)制，這仍然需要占用大量的GPU內(nèi)存；第三，Adam等優(yōu)化器的參數(shù)在訓(xùn)練過程中占用了大量GPU內(nèi)存，而混合精度訓(xùn)練使優(yōu)化器參數(shù)保持不變。從這個(gè)意義上說，需要更先進(jìn)的技術(shù)，如DeepSpeed的ZERO算法。

譯者介紹

朱先忠，51CTO社區(qū)編輯，51CTO專家博客、講師，濰坊一所高校計(jì)算機(jī)教師，自由編程界老兵一枚。

原文標(biāo)題：??The Mystery Behind the PyTorch Automatic Mixed Precision Library??，作者：Mengliu Zhao