前言

前幾日的文章，討論到了在訓練過程中，從資料的I/O、前處理到放入GPU之前的優化方式。今天將討論在GPU內的優化方式。

自動混和精度 - automatic mixed precision

我們在先前的訓練時，預設的狀況下，從讀取資料到模型的訓練以及驗證，全部的數值都是使用float32（單精度浮點數，可參考wiki）。

數值部份的資料格式，除了最常見的float32以外，其實還存在很多不同的種，例如：只有01的Boolean、整數常用的Int16以及今天要討論到的主角-----大小只有Float32一半的半精度浮點數Flost16。

• 大小只有一半的意思是，例如同樣是0這個數值，如果以float32儲存的話，大小是32bits；若是使用float16儲存，則只佔16bits
• float16的缺點則是因為容量較小，所以能描述的精準度較差，因此通常小數點的加減乘除起來的誤差也會較大，且值最大也只能容納到65519，超過則會被定義維為infinity。

所謂的AMP（automatic mixed precision，即自動混和精度），就是一種利用相同數值能以不同格式儲存所帶來的記憶體容量差異優勢來進行訓練的技術。更詳細的介紹可以觀看這篇文章，這裡截一張裡頭的重點圖：

其運作方式大致上是：

1. 以Float16進行forward
   • 中間的各層衍生物由原本的Float32改以Float16儲存，因此節省大量的記憶體
2. 透過loss以及backward計算出Float16梯度
3. 將梯度轉換回Float32
   • 這裡有滿多有趣的問題要處理，例如精準度不高，容易產生梯度消失，對應的作法會放大梯度
   • FLOAT16的最大值太小（只有65519），容易產生梯度爆炸，對應的作法則是會把梯度進行clip
4. 最後再以梯度更新回原始的float32格式，再進行權重的更新

其他一些加速的方式

除了今天主要介紹的AMP以外，還有一個個人覺得比較有感的加速則是Cudnn的Benchmark Mode。

詳細內容可以參考這篇討論，主要是由於cudnn內其實有多套不同的Convolution演算法實作，開啟這個功能主要是訓練時，cudnn會去自動調整針對當下tensor大小的最快最適合的演算法，來藉以加速。

實際實作很簡單，在pytroch裡面只要加上：

torch.backends.cudnn.benchmark = True

而Pytorch-Lightning則只要在Trainer內加上：

trainer = pl.Trainer(...,
                     benchmark = True)

就可以開啟這項功能了。

實際差異

現在讓我們實際來比較差異，這次主要針對有無AMP跟有無開啟CUDNN benchmark來做四組的比較：

左圖為訓練時的使用記憶體變化量，右圖是整體的計算時間

大概可以歸納出下列幾點：

• 純論GPU節省量來說，單開AMP最省
• 開啟CUDNN benchmark會比沒開來得快一些，但會多耗費一些記憶體（推估是為了記錄不同演算法的表現，所以需要額外資源）
• 開啟AMP反而計算比沒有開啟還要來得慢
  • 根據此篇，似乎在Volta框架後，nvidia有針對float32的計算特別硬體架構上的優化設計，因此已經加速到跟float16沒有什麼差別
  • 筆者使用的GPU是下一代ampere的3080TI，所以可能因此使用起來速度沒有太大差，反而多了float32、float16間的轉換而增加了計算時間

結語

• 介紹幾項優化GPU計算的方法
• 後續的實驗應該會採用 FLOAT32 + CUDNN BENCHMARK MODE 或是 AMP + 加大batch size的方式進行實驗