PyTorch 1.6 nightly增加了一個(gè)子模塊 amp，支持自動(dòng)混合精度訓(xùn)練。值得期待。來(lái)看看性能如何，相比Nvidia Apex 有哪些優(yōu)勢(shì)？

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

預(yù)計(jì)將在 PyTorch 1.6中推出的最令人興奮的附加功能之一是對(duì)自動(dòng)混合精度訓(xùn)練（automatic mixed-precision training）的支持。

混合精度訓(xùn)練是一種通過(guò)在半精度浮點(diǎn)數(shù) fp16上執(zhí)行盡可能多的操作來(lái)大幅度減少神經(jīng)網(wǎng)絡(luò)訓(xùn)練時(shí)間的技術(shù)，fp16 取代了PyTorch默認(rèn)的單精度浮點(diǎn)數(shù) fp32。最新一代 NVIDIA GPU 搭載了專門(mén)為快速 fp16矩陣運(yùn)算設(shè)計(jì)的特殊用途張量核（tensor cores）。

然而，到目前為止，這些張量核仍然很難用，因?yàn)樗枰謩?dòng)將精度降低的操作寫(xiě)入模型中。這就是自動(dòng)化混合精度訓(xùn)練的用武之地。即將發(fā)布的 torc h.cuda.amp API 將允許你只用五行代碼就可以在訓(xùn)練腳本中實(shí)現(xiàn)混合精度訓(xùn)練！

混合精度是如何工作的

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

在計(jì)算機(jī)工程中，像1.0151或566132.8這樣的十進(jìn)制數(shù)傳統(tǒng)上被表示為浮點(diǎn)數(shù)。由于我們可以有無(wú)限精確的數(shù)字(想象一下π) ，但存儲(chǔ)它們的空間是有限的，我們必須在精確度(在舍入數(shù)字前，我們可以在數(shù)字中包含的小數(shù)的數(shù)量)和大小(我們用來(lái)存儲(chǔ)數(shù)字的位數(shù))之間做出妥協(xié)。

浮點(diǎn)數(shù)的技術(shù)標(biāo)準(zhǔn) IEEE 754設(shè)定了以下標(biāo)準(zhǔn)：fp64, 又名雙精度或"double" ，最大舍入誤差 ~ 2^-52fp32, 又名單精度或"single"，最大舍入誤差 ~ 2 ^-23fp16, 又名半精度或"half" ，最大舍入誤差 ~ 2 ^-10。

Python float 類型為 fp64，而對(duì)內(nèi)存更敏感的PyTorch 使用 fp32作為默認(rèn)的 dtype。

混合精度訓(xùn)練的基本思想很簡(jiǎn)單: 精度減半(fp32→ fp16) ，訓(xùn)練時(shí)間減半。

最困難的是如何安全地做到這一點(diǎn)。

注意，浮點(diǎn)數(shù)越小，引起的舍入誤差就越大。對(duì)“足夠小“的浮點(diǎn)數(shù)執(zhí)行的任何操作都會(huì)將該值四舍五入到零！這就是所謂的underflowing，這是一個(gè)問(wèn)題，因?yàn)樵诜聪騻鞑ブ泻芏嗌踔链蠖鄶?shù)梯度更新值都非常小，但不為零。在反向傳播中舍入誤差累積可以把這些數(shù)字變成0或者 nans; 這會(huì)導(dǎo)致不準(zhǔn)確的梯度更新，影響你的網(wǎng)絡(luò)收斂。

2018年ICLR論文 Mixed Precision Training 發(fā)現(xiàn)，簡(jiǎn)單的在每個(gè)地方使用 fp16 會(huì)“吞掉”梯度更新小于2^-24的值—;—;大約占他們的示例網(wǎng)絡(luò)所有梯度更新的5% :

混合精度訓(xùn)練是一套技術(shù)，它允許你使用 fp16，而不會(huì)導(dǎo)致你的模型訓(xùn)練發(fā)生發(fā)散。這是三種不同技術(shù)的結(jié)合。

第一，維護(hù)兩個(gè)權(quán)重矩陣的副本，一個(gè)“主副本”用 fp32，一個(gè)半精度副本用 fp16。梯度更新使用 fp16矩陣計(jì)算，但更新于 fp32矩陣。這使得應(yīng)用梯度更新更加安全。

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

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

將這三種技術(shù)結(jié)合在一起，作者可以在顯著加速的時(shí)間內(nèi)訓(xùn)練好多種網(wǎng)絡(luò)以達(dá)到收斂。至于benchmarks，我建議讀一讀這篇只有9頁(yè)的論文！

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

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

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

更大的 fp16 矩陣乘法操作可以使用這個(gè)操作作為他們的基本構(gòu)件來(lái)實(shí)現(xiàn)。由于大多數(shù)反向傳播都可以歸結(jié)為矩陣乘法，張量核適用于網(wǎng)絡(luò)中幾乎任何計(jì)算密集層。

陷阱: 輸入矩陣必須是 fp16。 如果你正在使用帶有張量核的 GPU 進(jìn)行訓(xùn)練，而沒(méi)有使用混合精度訓(xùn)練，你不可能從你的顯卡中得到100% 的回報(bào)！ 在 fp32中定義的標(biāo)準(zhǔn) PyTorch 模型永遠(yuǎn)不會(huì)將任何 fp16數(shù)學(xué)運(yùn)算應(yīng)用到芯片上，因此所有這些極其強(qiáng)悍的張量核都將處于空閑狀態(tài)。

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

另一個(gè)值得記住的難題是firmware。盡管 CUDA 7.0或更高版本都支持張量核操作，但早期的實(shí)現(xiàn)據(jù)說(shuō)有很多 bug，所以使用 CUDA 10.0或更高版本很重要。

Pytorch 自動(dòng)混合精度是如何工作的

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

混合精度訓(xùn)練在技術(shù)上已經(jīng)永遠(yuǎn)成為可能: 手動(dòng)運(yùn)行部分網(wǎng)絡(luò)在 fp16中，并自己實(shí)現(xiàn)損失縮放。自動(dòng)混合精度訓(xùn)練中令人興奮的是“自動(dòng)”部分。只需要學(xué)習(xí)幾個(gè)新的 API 基本類型: torch.cuda.amp.GradScalar 和 torch.cuda.amp.autocast。啟用混合精度訓(xùn)練就像在你的訓(xùn)練腳本中插入正確的位置一樣簡(jiǎn)單！

為了演示，下面是使用混合精度訓(xùn)練的網(wǎng)絡(luò)訓(xùn)練循環(huán)的一段代碼。# NEW標(biāo)記定位了增加了新代碼的地方。

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 對(duì)象是 PyTorch 實(shí)現(xiàn)的損失縮放。回想一下在“混合精度如何工作”一節(jié)中提到，在訓(xùn)練期間，為了防止梯度變小到0，某種形式的縮放是必要的。最佳的損失乘數(shù)得足夠高以保留非常小的梯度，同時(shí)不能太高以至于導(dǎo)致非常大的梯度四舍五入到 inf產(chǎn)生相反的問(wèn)題。

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

通過(guò)這種方式逐級(jí)上下移動(dòng)損失乘數(shù)，PyTorch 可以隨著時(shí)間的推移近似得到合適的損失乘數(shù)。熟悉 TCP 擁塞控制的讀者應(yīng)該會(huì)發(fā)現(xiàn)這里的核心思想非常熟悉！該算法使用的準(zhǔn)確數(shù)字是可配置的，你可以直接從docstring中看到默認(rèn)值:

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

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

值得注意的是，GradScalar 可以檢測(cè)并停止overflows(因?yàn)?inf 總是壞的) ，但是它無(wú)法檢測(cè)和停止underflows(因?yàn)?通常是一個(gè)合法值)。如果你選擇的初始值太低，增長(zhǎng)間隔太長(zhǎng)，你的網(wǎng)絡(luò)可能會(huì)在 GradScalar 介入之前underflow并發(fā)散。由于這個(gè)原因，選擇一個(gè)非常大的初始值可能是一個(gè)好主意。

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

自動(dòng)混合精度訓(xùn)練拼圖的另一半是 torch.cuda.amp.autocast 上下文管理器。Autocast實(shí)現(xiàn)了 fp32-> fp16轉(zhuǎn)換?；叵胍幌隆盎旌暇仁侨绾喂ぷ鞯摹爸械膬?nèi)容，由于不同的操作以不同的速率累積誤差，并非所有的操作都可以在 fp16中安全運(yùn)行。下面的截圖來(lái)自 amp 模塊文檔，介紹了autocast如何處理 PyTorch 中可用的各種操作:

這個(gè)列表主要由矩陣乘法和卷積兩部分組成，還有簡(jiǎn)單的線性函數(shù)。

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

對(duì)數(shù)、指數(shù)、三角函數(shù)、正規(guī)函數(shù)、離散函數(shù)和(大)和在 fp16中是不安全的，必須在 fp32中執(zhí)行。

通過(guò)瀏覽這個(gè)列表，在我看來(lái)，大多數(shù)層都會(huì)從autocasting中受益，這要?dú)w功于它們內(nèi)部對(duì)基本線性代數(shù)操作的依賴，但大多數(shù)激活函數(shù)卻不是。卷積層是最大贏家。

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

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

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

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

Benchmarks性能

此時(shí)，我們已經(jīng)了解了什么是混合精度，什么是張量核，以及 PyTorch API 如何實(shí)現(xiàn)自動(dòng)混合精度。唯一剩下的就是看看一些真實(shí)世界的性能benchmarks！

我曾經(jīng)用自動(dòng)混合精度訓(xùn)練過(guò)三個(gè)非常不一樣的神經(jīng)網(wǎng)絡(luò)，還有一次沒(méi)用，通過(guò) Spell API 調(diào)用 V100s (上一代張量核)和 T4s (當(dāng)代張量核)。我分別使用了 AWS EC2實(shí)例、 p3.2xlarge 和 g4dn.xlarge，最近的 PyTorch 1.6 nightly 和 CUDA 10.0。所有模型的收斂都是一致的，即沒(méi)有一個(gè)模型發(fā)現(xiàn)混合精度網(wǎng)絡(luò)和原網(wǎng)絡(luò)在訓(xùn)練損失上有任何差異。訓(xùn)練的網(wǎng)絡(luò)如下:

前饋, 一個(gè)前饋神經(jīng)網(wǎng)絡(luò)，訓(xùn)練數(shù)據(jù)來(lái)自Kaggle比賽Rossman Store Samples UNet, 一個(gè)中等大小的原版UNet 圖像分割網(wǎng)絡(luò)， 在數(shù)據(jù)集Segmented Bob Ross Images 上訓(xùn)練 BERT, 一個(gè)大的 NLP transformer 模型，使用bert-base-uncased 骨干(通過(guò) huggingface)，及數(shù)據(jù)來(lái)自Kaggle競(jìng)賽 Twitter Sentiment Extraction

結(jié)果如下:

由于前饋網(wǎng)絡(luò)非常小，混合精度訓(xùn)練對(duì)它沒(méi)有任何好處。

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

BERT 是一個(gè)很大的模型，在這里使用混合精度訓(xùn)練節(jié)省時(shí)間，從中等模型的“很好”到了“必須擁有”。在Volta或Turing GPU 上訓(xùn)練，自動(dòng)混合精度將為大型模型減少50% 到60% 的訓(xùn)練時(shí)間！

這是一個(gè)巨大的優(yōu)勢(shì)，尤其是當(dāng)你考慮到增加的復(fù)雜性極小時(shí)—;—;只需要對(duì)模型訓(xùn)練腳本進(jìn)行四到五行代碼修改。在我看來(lái):

混合精度應(yīng)該是你對(duì)模型訓(xùn)練腳本進(jìn)行的最先性能優(yōu)化之一。

內(nèi)存呢？

正如我在“混合精度是如何工作的”一節(jié)中解釋的那樣，在內(nèi)存中fp16矩陣的大小是fp32矩陣的一半，因此，混合精度訓(xùn)練的另一個(gè)據(jù)稱的優(yōu)勢(shì)是內(nèi)存使用率。GPU 內(nèi)存的瓶頸遠(yuǎn)小于 GPU 的計(jì)算能力，但仍有很大的優(yōu)化價(jià)值。你的內(nèi)存使用效率越高，你可以在 GPU 上使用的batch size就越大。

PyTorch 在模型訓(xùn)練過(guò)程開(kāi)始時(shí)保留一定數(shù)量的 GPU 內(nèi)存，并在訓(xùn)練期間保留這些內(nèi)存。這可以防止其它進(jìn)程在訓(xùn)練過(guò)程中搶占過(guò)多的 GPU 內(nèi)存，迫使 PyTorch 訓(xùn)練腳本崩潰并出現(xiàn) OOM 錯(cuò)誤。

以下是啟用混合精度訓(xùn)練對(duì) PyTorch 內(nèi)存保留行為的影響:

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

總結(jié)

在即將發(fā)布的 PyTorch 1.6版本中，自動(dòng)混合精度訓(xùn)練是一個(gè)易于使用且功能強(qiáng)大的新特性，該版本承諾將在最新的 NVIDIA GPU 上運(yùn)行的大型模型訓(xùn)練工作加快60% 。

雖然這種技術(shù)已經(jīng)存在了一段時(shí)間，但是對(duì)于普通用戶來(lái)說(shuō)還不是很容易理解，因?yàn)橹钡浆F(xiàn)在它還沒(méi)有一個(gè)原生 PyTorch API。

要直接從源代碼中了解更多關(guān)于混合精度訓(xùn)練的信息，請(qǐng)參閱 PyTorch master 文檔中的automatic mixed precision package和automatic mixed precision examples頁(yè)面。

想自己測(cè)試一下這個(gè)功能？安裝最新的 PyTorch nightly非常簡(jiǎn)單: 查看 PyTorch 主頁(yè)上的說(shuō)明了解如何安裝。

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