[[333418]]

PyTorch 1.6 nightly增加了一個子模塊 amp，支持自動混合精度訓練。值得期待。來看看性能如何，相比Nvidia Apex 有哪些優勢？

即將在 PyTorch 1.6上發布的 torch.cuda.amp 混合精度訓練模塊實現了它的承諾，只需增加幾行新代碼就可以提高大型模型訓練50-60% 的速度。

預計將在 PyTorch 1.6中推出的最令人興奮的附加功能之一是對自動混合精度訓練（automatic mixed-precision training）的支持。

混合精度訓練是一種通過在半精度浮點數 fp16上執行盡可能多的操作來大幅度減少神經網絡訓練時間的技術，fp16 取代了PyTorch默認的單精度浮點數 fp32。最新一代 NVIDIA GPU 搭載了專門為快速 fp16矩陣運算設計的特殊用途張量核（tensor cores）。

然而，到目前為止，這些張量核仍然很難用，因為它需要手動將精度降低的操作寫入模型中。這就是自動化混合精度訓練的用武之地。即將發布的 torc h.cuda.amp API 將允許你只用五行代碼就可以在訓練腳本中實現混合精度訓練！

混合精度是如何工作的

在我們理解混合精度訓練是如何工作的之前，首先需要回顧一下浮點數。

在計算機工程中，像1.0151或566132.8這樣的十進制數傳統上被表示為浮點數。由于我們可以有無限精確的數字(想象一下π) ，但存儲它們的空間是有限的，我們必須在精確度(在舍入數字前，我們可以在數字中包含的小數的數量)和大小(我們用來存儲數字的位數)之間做出妥協。

浮點數的技術標準 IEEE 754設定了以下標準：fp64, 又名雙精度或"double" ，最大舍入誤差 ~ 2^-52fp32, 又名單精度或"single"，最大舍入誤差 ~ 2 ^-23fp16, 又名半精度或"half" ，最大舍入誤差 ~ 2 ^-10。

Python float 類型為 fp64，而對內存更敏感的PyTorch 使用 fp32作為默認的 dtype。

混合精度訓練的基本思想很簡單: 精度減半(fp32→ fp16) ，訓練時間減半。

最困難的是如何安全地做到這一點。

注意，浮點數越小，引起的舍入誤差就越大。對“足夠小“的浮點數執行的任何操作都會將該值四舍五入到零！這就是所謂的underflowing，這是一個問題，因為在反向傳播中很多甚至大多數梯度更新值都非常小，但不為零。在反向傳播中舍入誤差累積可以把這些數字變成0或者 nans; 這會導致不準確的梯度更新，影響你的網絡收斂。

2018年ICLR論文 Mixed Precision Training 發現，簡單的在每個地方使用 fp16 會“吞掉”梯度更新小于2^-24的值——大約占他們的示例網絡所有梯度更新的5% :

混合精度訓練是一套技術，它允許你使用 fp16，而不會導致你的模型訓練發生發散。這是三種不同技術的結合。

第一，維護兩個權重矩陣的副本，一個“主副本”用 fp32，一個半精度副本用 fp16。梯度更新使用 fp16矩陣計算，但更新于 fp32矩陣。這使得應用梯度更新更加安全。

第二，不同的向量操作以不同的速度累積誤差，因此要區別對待它們。有些操作在 fp16中總是安全的，而其它操作只在 fp32中是可靠的。與其用 fp16跑整個神經網絡，不如一些用半精度另外的用單精度。這種 dtypes 的混合就是為什么這種技術被稱為“混合精度”。

第三，使用損失縮放。損失縮放是指在執行反向傳播之前，將損失函數的輸出乘以某個標量數(論文建議從8開始)。乘性增加的損失值產生乘性增加的梯度更新值，“提升”許多梯度更新值到超過fp16的安全閾值2^-24。只要確保在應用梯度更新之前撤消縮放，并且不要選擇一個太大的縮放以至于產生 inf 權重更新(overflowing) ，從而導致網絡向相反的方向發散。

將這三種技術結合在一起，作者可以在顯著加速的時間內訓練好多種網絡以達到收斂。至于benchmarks，我建議讀一讀這篇只有9頁的論文！

張量核（tensor cores）是如何工作的

雖然混合精度訓練節省內存(fp16矩陣只有 fp32矩陣的一半大小) ，但如果沒有特殊的 GPU 支持，它并不能加速模型訓練。芯片上需要有可以加速半精度操作的東西。在最近幾代 NVIDIA GPU中這東西叫: 張量核。

張量核是一種新型的處理單元，針對一個非常特殊的操作進行了優化: 將兩個4 × 4 fp16矩陣相乘，然后將結果加到第三個4 × 4 fp16或 fp32矩陣(一個“融合乘法加（fused multiply add）”)中。

更大的 fp16 矩陣乘法操作可以使用這個操作作為他們的基本構件來實現。由于大多數反向傳播都可以歸結為矩陣乘法，張量核適用于網絡中幾乎任何計算密集層。

陷阱: 輸入矩陣必須是 fp16。 如果你正在使用帶有張量核的 GPU 進行訓練，而沒有使用混合精度訓練，你不可能從你的顯卡中得到100% 的回報！ 在 fp32中定義的標準 PyTorch 模型永遠不會將任何 fp16數學運算應用到芯片上，因此所有這些極其強悍的張量核都將處于空閑狀態。

張量核在2017年末在上一代Volta體系結構中被引入，當代Turing有了一些改進，并將在即將推出的Ampere中看到進一步的改進。云上通?？捎玫膬煽頖PU 是 V100(5120個 CUDA 核，600個張量核)和 T4(2560個 CUDA 核，320個張量核)。

另一個值得記住的難題是firmware。盡管 CUDA 7.0或更高版本都支持張量核操作，但早期的實現據說有很多 bug，所以使用 CUDA 10.0或更高版本很重要。

Pytorch 自動混合精度是如何工作的

有了這些重要的背景知識，我們終于可以開始深入研究新的 PyTorch amp API 了。

混合精度訓練在技術上已經永遠成為可能: 手動運行部分網絡在 fp16中，并自己實現損失縮放。自動混合精度訓練中令人興奮的是“自動”部分。只需要學習幾個新的 API 基本類型: torch.cuda.amp.GradScalar 和 torch.cuda.amp.autocast。啟用混合精度訓練就像在你的訓練腳本中插入正確的位置一樣簡單！

為了演示，下面是使用混合精度訓練的網絡訓練循環的一段代碼。# NEW標記定位了增加了新代碼的地方。

self.train() 
X = torch.tensor(X, dtype=torch.float32) 
y = torch.tensor(y, dtype=torch.float32) 
 
 
optimizer = torch.optim.Adam(self.parameters(), lr=self.max_lr) 
scheduler = torch.optim.lr_scheduler.OneCycleLR( 
    optimizer, self.max_lr, 
    cycle_momentum=False, 
    epochs=self.n_epochs, 
    steps_per_epoch=int(np.ceil(len(X) / self.batch_size)), 
) 
batches = torch.utils.data.DataLoader( 
    torch.utils.data.TensorDataset(X, y), 
    batch_size=self.batch_size, shuffle=True 
) 
 
 
# NEW 
scaler = torch.cuda.amp.GradScaler() 
 
 
for epoch in range(self.n_epochs): 
    for i, (X_batch, y_batch) in enumerate(batches): 
        X_batch = X_batch.cuda() 
        y_batch = y_batch.cuda() 
        optimizer.zero_grad() 
 
 
        # NEW 
        with torch.cuda.amp.autocast(): 
            y_pred = model(X_batch).squeeze() 
            loss = self.loss_fn(y_pred, y_batch) 
 
 
        # NEW 
        scaler.scale(loss).backward() 
        lv = loss.detach().cpu().numpy() 
        if i % 100 == 0: 
            print(f"Epoch {epoch + 1}/{self.n_epochs}; Batch {i}; Loss {lv}") 
 
 
        # NEW 
        scaler.step(optimizer) 
        scaler.update() 
 
        scheduler.s 

新的 PyTorch GradScaler 對象是 PyTorch 實現的損失縮放。回想一下在“混合精度如何工作”一節中提到，在訓練期間，為了防止梯度變小到0，某種形式的縮放是必要的。最佳的損失乘數得足夠高以保留非常小的梯度，同時不能太高以至于導致非常大的梯度四舍五入到 inf產生相反的問題。

PyTorch使用指數退避（exponential backoff）來解決這個問題。Gradscalar 以一個小的損失乘數開始，這個乘數每次會翻倍。這種逐漸加倍的行為一直持續到 GradScalar 遇到包含 inf 值的梯度更新。Gradscalar 丟棄這批數據(例如跳過梯度更新) ，將損失乘數減半，并重置其倍增時間。

通過這種方式逐級上下移動損失乘數，PyTorch 可以隨著時間的推移近似得到合適的損失乘數。熟悉 TCP 擁塞控制的讀者應該會發現這里的核心思想非常熟悉！該算法使用的準確數字是可配置的，你可以直接從docstring中看到默認值:

torch.cuda.amp.GradScaler( 
    init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, 
    growth_interval=2000, enabled=True 
)  

Gradscalar 需要對梯度更新計算(檢查是否溢出)和優化器(將丟棄的batches轉換為 no-op)進行控制，以實現其操作。這就是為什么 loss.backwards()被 scaler.scale(loss).backwards()取代， 以及 optimizer.step()被 scaler.step(optimizer)替換的原因。

值得注意的是，GradScalar 可以檢測并停止overflows(因為 inf 總是壞的) ，但是它無法檢測和停止underflows(因為0通常是一個合法值)。如果你選擇的初始值太低，增長間隔太長，你的網絡可能會在 GradScalar 介入之前underflow并發散。由于這個原因，選擇一個非常大的初始值可能是一個好主意。

最后，注意 GradScalar 是一個有狀態對象。使用此功能保存模型checkpoint需要和模型權重一起寫入和讀取磁盤。用 state _ dict 和 load _ state _ dict 對象方法(在 PyTorch 文檔中有介紹)可以很容易地做到這一點。

自動混合精度訓練拼圖的另一半是 torch.cuda.amp.autocast 上下文管理器。Autocast實現了 fp32-> fp16轉換?；叵胍幌?ldquo;混合精度是如何工作的“中的內容，由于不同的操作以不同的速率累積誤差，并非所有的操作都可以在 fp16中安全運行。下面的截圖來自 amp 模塊文檔，介紹了autocast如何處理 PyTorch 中可用的各種操作:

這個列表主要由矩陣乘法和卷積兩部分組成，還有簡單的線性函數。

這些操作在 fp16中是安全的，但是在輸入有 fp16和 fp32混合的情況下，這些操作具有向上適配(up-casting)規則，以確保它們不會出問題。注意，這個列表還包括另外兩個基本的線性代數運算: 矩陣/向量點積和向量叉積。

對數、指數、三角函數、正規函數、離散函數和(大)和在 fp16中是不安全的，必須在 fp32中執行。

通過瀏覽這個列表，在我看來，大多數層都會從autocasting中受益，這要歸功于它們內部對基本線性代數操作的依賴，但大多數激活函數卻不是。卷積層是最大贏家。

啟用sutocasting非常簡單。你只需要做的就是使用autocast上下文管理器包好模型的正向傳播:

with torch.cuda.amp.autocast(): 
    y_pred = model(X_batch).squeeze() 
    loss = self.loss_fn(y_pred, y_batch) 

以這種方式包裝前向傳播，可以自動打開后傳(如 loss.backwards ())的autocasting，因此不需要調用兩次autocast。

只要你遵循PyTorch 的最佳實踐(例如，避免in-place操作) ，autocasting基本上就可以“正常工作”。它甚至可以使用多GPU DistributedDataParallel API (只要遵循建議的策略，每個 GPU 只使用一個進程)。只需一個小調整，多GPU DataParallel API也可以用。Pytorch 文檔中的 Automatic Mixed Precision Examples 頁面的“Working with multiple GPUs”部分是關于這個主題的一個方便的參考。個人觀點，有一個要記住的重點是: "優先用 binary cross entropy with logits 而不是 binary cross entropy"。

Benchmarks性能

此時，我們已經了解了什么是混合精度，什么是張量核，以及 PyTorch API 如何實現自動混合精度。唯一剩下的就是看看一些真實世界的性能benchmarks！

我曾經用自動混合精度訓練過三個非常不一樣的神經網絡，還有一次沒用，通過 Spell API 調用 V100s (上一代張量核)和 T4s (當代張量核)。我分別使用了 AWS EC2實例、 p3.2xlarge 和 g4dn.xlarge，最近的 PyTorch 1.6 nightly 和 CUDA 10.0。所有模型的收斂都是一致的，即沒有一個模型發現混合精度網絡和原網絡在訓練損失上有任何差異。訓練的網絡如下:

• 前饋, 一個前饋神經網絡，訓練數據來自Kaggle比賽Rossman Store Samples
• UNet, 一個中等大小的原版UNet 圖像分割網絡， 在數據集Segmented Bob Ross Images 上訓練
• BERT, 一個大的 NLP transformer 模型，使用bert-base-uncased 骨干(通過 huggingface)，及數據來自Kaggle競賽 Twitter Sentiment Extraction

結果如下:

由于前饋網絡非常小，混合精度訓練對它沒有任何好處。

UNet 是一個中等規模的卷積模型，共有7,703,497個參數，從混合精度訓練中得到了顯著的好處。有趣的是，雖然 V100和 T4都受益于混合精度訓練，但 T4的好處要大得多: 節省5%時間vs. 高達30%的時間。

BERT 是一個很大的模型，在這里使用混合精度訓練節省時間，從中等模型的“很好”到了“必須擁有”。在Volta或Turing GPU 上訓練，自動混合精度將為大型模型減少50% 到60% 的訓練時間！

這是一個巨大的優勢，尤其是當你考慮到增加的復雜性極小時——只需要對模型訓練腳本進行四到五行代碼修改。在我看來:

混合精度應該是你對模型訓練腳本進行的最先性能優化之一。

內存呢？

正如我在“混合精度是如何工作的”一節中解釋的那樣，在內存中fp16矩陣的大小是fp32矩陣的一半，因此，混合精度訓練的另一個據稱的優勢是內存使用率。GPU 內存的瓶頸遠小于 GPU 的計算能力，但仍有很大的優化價值。你的內存使用效率越高，你可以在 GPU 上使用的batch size就越大。

PyTorch 在模型訓練過程開始時保留一定數量的 GPU 內存，并在訓練期間保留這些內存。這可以防止其它進程在訓練過程中搶占過多的 GPU 內存，迫使 PyTorch 訓練腳本崩潰并出現 OOM 錯誤。

以下是啟用混合精度訓練對 PyTorch 內存保留行為的影響:

有趣的是，雖然兩個較大的模型都看到了切換到混合精度的好處，UNet 從切換中得到的好處比 BERT 多得多。PyTorch 內存分配行為對我來說非常不透明，所以我不知道為什么會出現這種情況。

總結

在即將發布的 PyTorch 1.6版本中，自動混合精度訓練是一個易于使用且功能強大的新特性，該版本承諾將在最新的 NVIDIA GPU 上運行的大型模型訓練工作加快60% 。

雖然這種技術已經存在了一段時間，但是對于普通用戶來說還不是很容易理解，因為直到現在它還沒有一個原生 PyTorch API。

要直接從源代碼中了解更多關于混合精度訓練的信息，請參閱 PyTorch master 文檔中的automatic mixed precision package和automatic mixed precision examples頁面。

想自己測試一下這個功能？安裝最新的 PyTorch nightly非常簡單: 查看 PyTorch 主頁上的說明了解如何安裝。

想要自己復現這些benchmarks嗎？所有模型源代碼都可以在 GitHub 上的 ResidentMario/spell-feedforward-rossman, ResidentMario/spell-unet-bob-ross, 和 ResidentMario/spell-tweet-sentiment-extraction 庫中獲得。