PyTorch 1.6 nightly增加了一個子模塊 amp，支持自動混合精度訓練。值得期待。來看看性能如何，相比Nvidia Apex 有哪些優(yōu)勢？
即將在 PyTorch 1.6上發(fā)布的 torch.cuda.amp 混合精度訓練模塊實現(xiàn)了它的承諾，只需增加幾行新代碼就可以提高大型模型訓練50-60% 的速度。
預計將在 PyTorch 1.6中推出的最令人興奮的附加功能之一是對自動混合精度訓練（automatic mixed-precision training）的支持。
混合精度訓練是一種通過在半精度浮點數(shù) fp16上執(zhí)行盡可能多的操作來大幅度減少神經(jīng)網(wǎng)絡訓練時間的技術(shù)，fp16 取代了PyTorch默認的單精度浮點數(shù) fp32。最新一代 NVIDIA GPU 搭載了專門為快速 fp16矩陣運算設計的特殊用途張量核（tensor cores）。
然而，到目前為止，這些張量核仍然很難用，因為它需要手動將精度降低的操作寫入模型中。這就是自動化混合精度訓練的用武之地。即將發(fā)布的 torc h.cuda.amp API 將允許你只用五行代碼就可以在訓練腳本中實現(xiàn)混合精度訓練！
混合精度是如何工作的
在我們理解混合精度訓練是如何工作的之前，首先需要回顧一下浮點數(shù)。
在計算機工程中，像1.0151或566132.8這樣的十進制數(shù)傳統(tǒng)上被表示為浮點數(shù)。由于我們可以有無限精確的數(shù)字(想象一下π) ，但存儲它們的空間是有限的，我們必須在精確度(在舍入數(shù)字前，我們可以在數(shù)字中包含的小數(shù)的數(shù)量)和大小(我們用來存儲數(shù)字的位數(shù))之間做出妥協(xié)。
浮點數(shù)的技術(shù)標準 IEEE 754設定了以下標準：fp64, 又名雙精度或double ，舍入誤差 ~ 2^-52fp32, 又名單精度或single，舍入誤差 ~ 2 ^-23fp16, 又名半精度或half ，舍入誤差 ~ 2 ^-10。
Python float 類型為 fp64，而對內(nèi)存更敏感的PyTorch 使用 fp32作為默認的 dtype。
混合精度訓練的基本思想很簡單: 精度減半(fp32→ fp16) ，訓練時間減半。
最困難的是如何安全地做到這一點。
注意，浮點數(shù)越小，引起的舍入誤差就越大。對“足夠小“的浮點數(shù)執(zhí)行的任何操作都會將該值四舍五入到零！這就是所謂的underflowing，這是一個問題，因為在反向傳播中很多甚至大多數(shù)梯度更新值都非常小，但不為零。在反向傳播中舍入誤差累積可以把這些數(shù)字變成0或者 nans; 這會導致不準確的梯度更新，影響你的網(wǎng)絡收斂。
2018年ICLR論文 Mixed Precision Training 發(fā)現(xiàn)，簡單的在每個地方使用 fp16 會“吞掉”梯度更新小于2^-24的值——大約占他們的示例網(wǎng)絡所有梯度更新的5% :
混合精度訓練是一套技術(shù)，它允許你使用 fp16，而不會導致你的模型訓練發(fā)生發(fā)散。這是三種不同技術(shù)的結(jié)合。
第一，維護兩個權(quán)重矩陣的副本，一個“主副本”用 fp32，一個半精度副本用 fp16。梯度更新使用 fp16矩陣計算，但更新于 fp32矩陣。這使得應用梯度更新更加安全。
第二，不同的向量操作以不同的速度累積誤差，因此要區(qū)別對待它們。有些操作在 fp16中總是安全的，而其它操作只在 fp32中是可靠的。與其用 fp16跑整個神經(jīng)網(wǎng)絡，不如一些用半精度另外的用單精度。這種 dtypes 的混合就是為什么這種技術(shù)被稱為“混合精度”。
第三，使用損失縮放。損失縮放是指在執(zhí)行反向傳播之前，將損失函數(shù)的輸出乘以某個標量數(shù)(論文建議從8開始)。乘性增加的損失值產(chǎn)生乘性增加的梯度更新值，“提升”許多梯度更新值到超過fp16的安全閾值2^-24。只要確保在應用梯度更新之前撤消縮放，并且不要選擇一個太大的縮放以至于產(chǎn)生 inf 權(quán)重更新(overflowing) ，從而導致網(wǎng)絡向相反的方向發(fā)散。
將這三種技術(shù)結(jié)合在一起，作者可以在顯著加速的時間內(nèi)訓練好多種網(wǎng)絡以達到收斂。至于benchmarks，我建議讀一讀這篇只有9頁的論文！
張量核（tensor cores）是如何工作的
雖然混合精度訓練節(jié)省內(nèi)存(fp16矩陣只有 fp32矩陣的一半大小) ，但如果沒有特殊的 GPU 支持，它并不能加速模型訓練。芯片上需要有可以加速半精度操作的東西。在最近幾代 NVIDIA GPU中這東西叫: 張量核。
張量核是一種新型的處理單元，針對一個非常特殊的操作進行了優(yōu)化: 將兩個4 × 4 fp16矩陣相乘，然后將結(jié)果加到第三個4 × 4 fp16或 fp32矩陣(一個“融合乘法加（fused multiply add）”)中。
更大的 fp16 矩陣乘法操作可以使用這個操作作為他們的基本構(gòu)件來實現(xiàn)。由于大多數(shù)反向傳播都可以歸結(jié)為矩陣乘法，張量核適用于網(wǎng)絡中幾乎任何計算密集層。
陷阱: 輸入矩陣必須是 fp16。 如果你正在使用帶有張量核的 GPU 進行訓練，而沒有使用混合精度訓練，你不可能從你的顯卡中得到100% 的回報！ 在 fp32中定義的標準 PyTorch 模型永遠不會將任何 fp16數(shù)學運算應用到芯片上，因此所有這些極其強悍的張量核都將處于空閑狀態(tài)。
張量核在2017年末在上一代Volta體系結(jié)構(gòu)中被引入，當代Turing有了一些改進，并將在即將推出的Ampere中看到進一步的改進。云上通常可用的兩款GPU 是 V100(5120個 CUDA 核，600個張量核)和 T4(2560個 CUDA 核，320個張量核)。
另一個值得記住的難題是firmware。盡管 CUDA 7.0或更高版本都支持張量核操作，但早期的實現(xiàn)據(jù)說有很多 bug，所以使用 CUDA 10.0或更高版本很重要。
Pytorch 自動混合精度是如何工作的
有了這些重要的背景知識，我們終于可以開始深入研究新的 PyTorch amp API 了。
混合精度訓練在技術(shù)上已經(jīng)永遠成為可能: 手動運行部分網(wǎng)絡在 fp16中，并自己實現(xiàn)損失縮放。自動混合精度訓練中令人興奮的是“自動”部分。只需要學習幾個新的 API 基本類型: torch.cuda.amp.GradScalar 和 torch.cuda.amp.autocast。啟用混合精度訓練就像在你的訓練腳本中插入正確的位置一樣簡單！
為了演示，下面是使用混合精度訓練的網(wǎng)絡訓練循環(huán)的一段代碼。# 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)#NEWscaler=torch.cuda.amp.GradScaler()forepochinrange(self.n_epochs):fori,(X_batch,y_batch)inenumerate(batches):X_batch=X_batch.cuda()y_batch=y_batch.cuda()optimizer.zero_grad()#NEWwithtorch.cuda.amp.autocast():y_pred=model(X_batch).squeeze()loss=self.loss_fn(y_pred,y_batch)#NEWscaler.scale(loss).backward()lv=loss.detach().cpu().numpy()ifi%100==0:print(fEpoch{epoch+1}/{self.n_epochs};Batch{i};Loss{lv})#NEWscaler.step(optimizer)scaler.update()scheduler.s
新的 PyTorch GradScaler 對象是 PyTorch 實現(xiàn)的損失縮放。回想一下在“混合精度如何工作”一節(jié)中提到，在訓練期間，為了防止梯度變小到0，某種形式的縮放是必要的。的損失乘數(shù)得足夠高以保留非常小的梯度，同時不能太高以至于導致非常大的梯度四舍五入到 inf產(chǎn)生相反的問題。
PyTorch使用指數(shù)退避（exponential backoff）來解決這個問題。Gradscalar 以一個小的損失乘數(shù)開始，這個乘數(shù)每次會翻倍。這種逐漸加倍的行為一直持續(xù)到 GradScalar 遇到包含 inf 值的梯度更新。Gradscalar 丟棄這批數(shù)據(jù)(例如跳過梯度更新) ，將損失乘數(shù)減半，并重置其倍增時間。
通過這種方式逐級上下移動損失乘數(shù)，PyTorch 可以隨著時間的推移近似得到合適的損失乘數(shù)。熟悉 TCP 擁塞控制的讀者應該會發(fā)現(xiàn)這里的核心思想非常熟悉！該算法使用的準確數(shù)字是可配置的，你可以直接從docstring中看到默認值:
torch.cuda.amp.GradScaler(init_scale=65536.0,growth_factor=2.0,backoff_factor=0.5,growth_interval=2000,enabled=True)
Gradscalar 需要對梯度更新計算(檢查是否溢出)和優(yōu)化器(將丟棄的batches轉(zhuǎn)換為 no-op)進行控制，以實現(xiàn)其操作。這就是為什么 loss.backwards()被 scaler.scale(loss).backwards()取代， 以及 optimizer.step()被 scaler.step(optimizer)替換的原因。
值得注意的是，GradScalar 可以檢測并停止overflows(因為 inf 總是壞的) ，但是它無法檢測和停止underflows(因為0通常是一個合法值)。如果你選擇的初始值太低，增長間隔太長，你的網(wǎng)絡可能會在 GradScalar 介入之前underflow并發(fā)散。由于這個原因，選擇一個非常大的初始值可能是一個好主意。
最后，注意 GradScalar 是一個有狀態(tài)對象。使用此功能保存模型checkpoint需要和模型權(quán)重一起寫入和讀取磁盤。用 state _ dict 和 load _ state _ dict 對象方法(在 PyTorch 文檔中有介紹)可以很容易地做到這一點。
自動混合精度訓練拼圖的另一半是 torch.cuda.amp.autocast 上下文管理器。Autocast實現(xiàn)了 fp32-> fp16轉(zhuǎn)換?；叵胍幌?ldquo;混合精度是如何工作的“中的內(nèi)容，由于不同的操作以不同的速率累積誤差，并非所有的操作都可以在 fp16中安全運行。下面的截圖來自 amp 模塊文檔，介紹了autocast如何處理 PyTorch 中可用的各種操作:
這個列表主要由矩陣乘法和卷積兩部分組成，還有簡單的線性函數(shù)。
這些操作在 fp16中是安全的，但是在輸入有 fp16和 fp32混合的情況下，這些操作具有向上適配(up-casting)規(guī)則，以確保它們不會出問題。注意，這個列表還包括另外兩個基本的線性代數(shù)運算: 矩陣/向量點積和向量叉積。
對數(shù)、指數(shù)、三角函數(shù)、正規(guī)函數(shù)、離散函數(shù)和(大)和在 fp16中是不安全的，必須在 fp32中執(zhí)行。
通過瀏覽這個列表，在我看來，大多數(shù)層都會從autocasting中受益，這要歸功于它們內(nèi)部對基本線性代數(shù)操作的依賴，但大多數(shù)激活函數(shù)卻不是。卷積層是贏家。
啟用sutocasting非常簡單。你只需要做的就是使用autocast上下文管理器包好模型的正向傳播:
withtorch.cuda.amp.autocast():y_pred=model(X_batch).squeeze()loss=self.loss_fn(y_pred,y_batch)
以這種方式包裝前向傳播，可以自動打開后傳(如 loss.backwards ())的autocasting，因此不需要調(diào)用兩次autocast。
只要你遵循PyTorch 的實踐(例如，避免in-place操作) ，autocasting基本上就可以“正常工作”。它甚至可以使用多GPU DistributedDataParallel API (只要遵循建議的策略，每個 GPU 只使用一個進程)。只需一個小調(diào)整，多GPU DataParallel API也可以用。Pytorch 文檔中的 Automatic Mixed Precision Examples 頁面的“Working with multiple GPUs”部分是關于這個主題的一個方便的參考。個人觀點，有一個要記住的重點是: 優(yōu)先用 binary cross entropy with logits 而不是 binary cross entropy。
Benchmarks性能
此時，我們已經(jīng)了解了什么是混合精度，什么是張量核，以及 PyTorch API 如何實現(xiàn)自動混合精度。唯一剩下的就是看看一些真實世界的性能benchmarks！
我曾經(jīng)用自動混合精度訓練過三個非常不一樣的神經(jīng)網(wǎng)絡，還有一次沒用，通過 Spell API 調(diào)用 V100s (上一代張量核)和 T4s (當代張量核)。我分別使用了 AWS EC2實例、 p3.2xlarge 和 g4dn.xlarge，最近的 PyTorch 1.6 nightly 和 CUDA 10.0。所有模型的收斂都是一致的，即沒有一個模型發(fā)現(xiàn)混合精度網(wǎng)絡和原網(wǎng)絡在訓練損失上有任何差異。訓練的網(wǎng)絡如下:
前饋, 一個前饋神經(jīng)網(wǎng)絡，訓練數(shù)據(jù)來自Kaggle比賽Rossman Store Samples UNet, 一個中等大小的原版UNet 圖像分割網(wǎng)絡， 在數(shù)據(jù)集Segmented Bob Ross Images 上訓練 BERT, 一個大的 NLP transformer 模型，使用bert-base-uncased 骨干(通過 huggingface)，及數(shù)據(jù)來自Kaggle競賽 Twitter Sentiment Extraction
結(jié)果如下:
由于前饋網(wǎng)絡非常小，混合精度訓練對它沒有任何好處。
UNet 是一個中等規(guī)模的卷積模型，共有7,703,497個參數(shù)，從混合精度訓練中得到了顯著的好處。有趣的是，雖然 V100和 T4都受益于混合精度訓練，但 T4的好處要大得多: 節(jié)省5%時間vs. 高達30%的時間。
BERT 是一個很大的模型，在這里使用混合精度訓練節(jié)省時間，從中等模型的“很好”到了“必須擁有”。在Volta或Turing GPU 上訓練，自動混合精度將為大型模型減少50% 到60% 的訓練時間！
這是一個巨大的優(yōu)勢，尤其是當你考慮到增加的復雜性極小時——只需要對模型訓練腳本進行四到五行代碼修改。在我看來:
混合精度應該是你對模型訓練腳本進行的最先性能優(yōu)化之一。
內(nèi)存呢？
正如我在“混合精度是如何工作的”一節(jié)中解釋的那樣，在內(nèi)存中fp16矩陣的大小是fp32矩陣的一半，因此，混合精度訓練的另一個據(jù)稱的優(yōu)勢是內(nèi)存使用率。GPU 內(nèi)存的瓶頸遠小于 GPU 的計算能力，但仍有很大的優(yōu)化價值。你的內(nèi)存使用效率越高，你可以在 GPU 上使用的batch size就越大。
PyTorch 在模型訓練過程開始時保留一定數(shù)量的 GPU 內(nèi)存，并在訓練期間保留這些內(nèi)存。這可以防止其它進程在訓練過程中搶占過多的 GPU 內(nèi)存，迫使 PyTorch 訓練腳本崩潰并出現(xiàn) OOM 錯誤。
以下是啟用混合精度訓練對 PyTorch 內(nèi)存保留行為的影響:
有趣的是，雖然兩個較大的模型都看到了切換到混合精度的好處，UNet 從切換中得到的好處比 BERT 多得多。PyTorch 內(nèi)存分配行為對我來說非常不透明，所以我不知道為什么會出現(xiàn)這種情況。
總結(jié)
在即將發(fā)布的 PyTorch 1.6版本中，自動混合精度訓練是一個易于使用且功能強大的新特性，該版本承諾將在最新的 NVIDIA GPU 上運行的大型模型訓練工作加快60% 。
雖然這種技術(shù)已經(jīng)存在了一段時間，但是對于普通用戶來說還不是很容易理解，因為直到現(xiàn)在它還沒有一個原生 PyTorch API。
要直接從源代碼中了解更多關于混合精度訓練的信息，請參閱 PyTorch master 文檔中的automatic mixed precision package和automatic mixed precision examples頁面。
想自己測試一下這個功能？安裝最新的 PyTorch nightly非常簡單: 查看 PyTorch 主頁上的說明了解如何安裝。
想要自己復現(xiàn)這些benchmarks嗎？所有模型源代碼都可以在 GitHub 上的 ResidentMario/spell-feedforward-rossman, ResidentMario/spell-unet-bob-ross, 和 ResidentMario/spell-tweet-sentiment-extraction 庫中獲得。
重慶分公司
重慶分公司
