【PyTorch】加速不香嗎？| 基于Apex的混合精度加速

點選藍字關注我們

掃碼關注我們

公衆号 : 計算機視覺戰隊

加入我們，大量論文代碼下載下傳連結

 背景 

我們提到圓周率 π 的時候，它有很多種表達方式，既可以用數學常數3.14159表示，也可以用一長串1和0的二進制長串表示。

圓周率 π 是個無理數，既小數位無限且不循環。是以，在使用圓周率進行計算時，人和計算機都必須根據精度需要将小數點後的數字四舍五入。

在國小的時候，國小生們可能隻會用手算的方式計算數學題目，圓周率的數值也隻能計算到小數點後兩位——3.14；而高中生使用圖形電腦可能會使圓周率數值排到小數點後10位，更加精确地表示圓周率。在計算機科學中，這被稱為精度，它通常以二進制數字來衡量，而非小數。

對于複雜的科學模拟，開發人員長期以來一直都依靠高精度數學來研究諸如宇宙大爆炸，或是預測數百萬個原子之間的互相作用。

數字位數越高，或是小數點後位數越多，意味着科學家可以在更大範圍内的數值内展現兩個數值的變化。借此，科學家可以對最大的星系，或是最小的粒子進行精确計算。

但是，計算精度越高，意味着所需的計算資源、資料傳輸和記憶體存儲就越多。其成本也會更大，同時也會消耗更多的功率。

由于并非每個工作負載都需要高精度，是以 AI 和 HPC 研究人員可以通過混合或比對不同級别的精度的方式進行運算，進而使效益最大化。NVIDIA Tensor Core GPU 支援多精度和混合精度技術，能夠讓開發者優化計算資源并加快 AI 應用程式及其推理功能的訓練。

單精度、雙精度和半精度浮點格式之間的差別

IEEE 浮點算術标準是用來衡量計算機上以二進制所表示數字精度的通用約定。在雙精度格式中，每個數字占用64位，單精度格式占用32位，而半精度僅16位。

要了解其中工作原理，我們可以拿圓周率舉例。在傳統科學記數法中，圓周率表示為3.14 x100。但是計算機将這些資訊以二進制形式存儲為浮點，即一系列的1和0，它們代表一個數字及其對應的指數，在這種情況下圓周率則表示為1.1001001 x 21。

在單精度32位格式中，1位用于訓示數字為正數還是負數。指數保留了8位，這是因為它為二進制，将2進到高位。其餘23位用于表示組成該數字的數字，稱為有效數字。

而在雙精度下，指數保留11位，有效位數為52位，進而極大地擴充了它可以表示的數字範圍和大小。半精度則是表示範圍更小，其指數隻有5位，有效位數隻有10位。

圓周率在每個精度級别表現如下：

多精度和混合精度計算的差異

多精度計算意味着使用能夠以不同精度進行計算的處理器，在需要使用高精度進行計算的部分使用雙精度，并在應用程式的其他部分使用半精度或單精度算法。

混合精度（也稱為超精度）計算則是在單個操作中使用不同的精度級别，進而在不犧牲精度的情況下實作計算效率。

在混合精度中，計算從半精度值開始，以進行快速矩陣數學運算。但是随着數字的計算，機器會以更高的精度存儲結果。例如，如果将兩個16位矩陣相乘，則結果為32位大小。

使用這種方法，在應用程式結束計算時，其累積得到結果，在準确度上可與使用雙精度算法運算得到的結果相媲美。

這項技術可以将傳統的雙精度應用程式加速多達25倍，同時減少了運作所需的記憶體、時間和功耗。它可用于 AI 和模拟 HPC 工作負載。

随着混合精度算法在現代超級計算應用程式中的普及，HPC 專家 Jack Dongarra 提出了一個新的基準，即 HPL-AI，以評估超級計算機在混合精度計算上的性能。當 NVIDIA 在全球最快的超級計算機 Summit 上運作 HPL-AI 計算時，該系統達到了前所未有的性能水準，接近550 petaflop，比其在 TOP500 榜單上的官方性能快了3倍。

 理論 

PyTorch實作

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 這裡是“歐一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
           

但是如果你希望對FP16和Apex有更深入的了解，或是在使用中遇到了各種不明是以的“Nan”的同學，可以接着讀下去，後面會有一些有趣的理論知識和瓦礫最近一個月使用Apex遇到的各種bug，不過當你深入了解并解決掉這些bug後，你就可以徹底擺脫“慢吞吞”的FP32啦。

理論部分

為了充分了解混合精度的原理，以及API的使用，先補充一點基礎的理論知識。

什麼是FP16？

半精度浮點數是一種計算機使用的二進制浮點數資料類型，使用2位元組（16位）存儲。

其中，sign位表示正負，exponent位表示指數（$2^{n-15+{1}(n=0)}$），fraction位表示的是分數（$\frac{m}{1024}$）。其中當指數為零的時候，下圖加号左邊為0，其他情況為1。

FP16的表示範例

為什麼需要FP16？

在使用FP16之前，我想再贅述一下為什麼我們使用FP16。

• 減少顯存占用

現在模型越來越大，當你使用Bert這一類的預訓練模型時，往往顯存就被模型及模型計算占去大半，當想要使用更大的Batch Size的時候會顯得捉襟見肘。由于FP16的記憶體占用隻有FP32的一半，自然地就可以幫助訓練過程節省一半的顯存空間。

• 加快訓練和推斷的計算

與普通的空間時間Trade-off的加速方法不同，FP16除了能節約記憶體，還能同時節省模型的訓練時間。在大部分的測試中，基于FP16的加速方法能夠給模型訓練帶來多一倍的加速體驗（爽感類似于兩倍速看肥皂劇）。

• 張量核心的普及

硬體的發展同樣也推動着模型計算的加速，随着Nvidia張量核心（Tensor Core）的普及，16bit計算也一步步走向成熟，低精度計算也是未來深度學習的一個重要趨勢，再不學習就out啦。

FP16帶來的問題：量化誤差

這個部分是整個部落格最重要的理論核心。講了這麼多FP16的好處，那麼使用FP16的時候有沒有什麼問題呢？當然有。FP16帶來的問題主要有兩個：1. 溢出錯誤；2. 舍入誤差。

• 溢出錯誤（Grad Overflow / Underflow）

由于FP16的動态範圍（$6 \times 10^{-8} \sim 65504$）比FP32的動态範圍（$1.4 \times 10^{-45} \sim 1.7 \times 10^{38}$）要狹窄很多，是以在計算過程中很容易出現上溢出（Overflow，$g>65504$）和下溢出（Underflow，$g<6\times10^{-8}$）的錯誤，溢出之後就會出現“Nan”的問題。

在深度學習中，由于激活函數的的梯度往往要比權重梯度小，更易出現下溢出的情況。

• 舍入誤差（Rounding Error）

舍入誤差指的是當梯度過小，小于目前區間内的最小間隔時，該次梯度更新可能會失敗，用一張圖清晰地表示：

解決問題的辦法：混合精度訓練+動态損失放大

• 混合精度訓練（Mixed Precision）

混合精度訓練的精髓在于“在記憶體中用FP16做儲存和乘法進而加速計算，用FP32做累加避免舍入誤差”。混合精度訓練的政策有效地緩解了舍入誤差的問題。

• 損失放大（Loss Scaling）

即使用了混合精度訓練，還是會存在無法收斂的情況，原因是激活梯度的值太小，造成了下溢出（Underflow）。損失放大的思路是：

1. 反向傳播前，将損失變化（dLoss）手動增大$2^k$倍，是以反向傳播時得到的中間變量（激活函數梯度）則不會溢出；
2. 反向傳播後，将權重梯度縮$2^k$倍，恢複正常值。

 Apex的新API：Automatic Mixed Precision (AMP)

曾經的Apex混合精度訓練的api仍然需要手動half模型已經輸入的資料，比較麻煩，現在新的api隻需要三行代碼即可無痛使用：

1 2 3 4

from apex import amp model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 這裡是“歐一”，不是“零一” with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()

1. opt_level

   其中隻有一個

   opt_level

   需要使用者自行配置：

• O0

  ：純FP32訓練，可以作為accuracy的baseline；

• O1

  ：混合精度訓練（推薦使用），根據黑白名單自動決定使用FP16（GEMM, 卷積）還是FP32（Softmax）進行計算。

• O2

  ：“幾乎FP16”混合精度訓練，不存在黑白名單，除了Batch norm，幾乎都是用FP16計算。

• O3

  ：純FP16訓練，很不穩定，但是可以作為speed的baseline；

• 動态損失放大（Dynamic Loss Scaling）

  AMP預設使用動态損失放大，為了充分利用FP16的範圍，緩解舍入誤差，盡量使用最高的放大倍數（$2^{24}$），如果産生了上溢出（Overflow），則跳過參數更新，縮小放大倍數使其不溢出，在一定步數後（比如2000步）會再嘗試使用大的scale來充分利用FP16的範圍：

• 

   踩過的坑 

  這一部分是整篇最幹貨的部分，在最近在apex使用中的踩過的所有的坑，由于apex報錯并不明顯，常常debug得讓人很沮喪，但隻要注意到以下的點，95%的情況都可以暢通無阻了：

  1. 判斷你的GPU是否支援FP16：構擁有Tensor Core的GPU（2080Ti、Titan、Tesla等），不支援的（Pascal系列）就不建議折騰了。
  2. 常數的範圍：為了保證計算不溢出，首先要保證人為設定的常數（包括調用的源碼中的）不溢出，如各種epsilon，INF等。
  3. Dimension最好是8的倍數：Nvidia官方的文檔的2.2條表示，次元都是8的倍數的時候，性能最好。
  4. 涉及到sum的操作要小心，很容易溢出，類似Softmax的操作建議用官方API，并定義成layer寫在模型初始化裡。
  5. 模型書寫要規範：自定義的Layer寫在模型初始化函數裡，graph計算寫在forward裡。
  6. 某些不常用的函數，在使用前需要注冊：amp.register_float_function(torch, 'sigmoid')
  7. 某些函數（如einsum）暫不支援FP16加速，建議不要用的太heavy，xlnet的實作改FP16困擾了我很久。
  8. 需要操作模型參數的子產品（類似EMA），要使用AMP封裝後的model。
  9. 需要操作梯度的子產品必須在optimizer的step裡，不然AMP不能判斷grad是否為Nan。

   總結 

  這篇從理論到實踐地介紹了混合精度計算以及Apex新API（AMP）的使用方法。現在在做深度學習模型的時候，幾乎都會第一時間把代碼改成混合精度訓練的了，速度快，精度還不減，确實是調參煉丹必備神器。目前網上還并沒有看到關于AMP以及使用時會遇到的坑的中文部落格，是以這一篇也是希望大家在使用的時候可以少花一點時間debug。

  

  如果想加入我們“計算機視覺研究院”，請掃二維碼加入我們。我們會按照你的需求将你拉入對應的學習群！

  計算機視覺研究院主要涉及深度學習領域，主要緻力于人臉檢測、人臉識别，多目标檢測、目标跟蹤、圖像分割等研究方向。研究院接下來會不斷分享最新的論文算法新架構，我們這次改革不同點就是，我們要着重”研究“。之後我們會針對相應領域分享實踐過程，讓大家真正體會擺脫理論的真實場景，培養愛動手程式設計愛動腦思考的習慣！

  

  掃碼關注我們

  公衆号 : 計算機視覺戰隊