機(jī)器之心編輯部
飛槳自動(dòng)混合精度技術(shù)，讓你的訓(xùn)練速度飛起來(lái)。
隨著生活節(jié)奏的加快，「等待」已經(jīng)越來(lái)越成為人們希望遠(yuǎn)離的事情。但是在深度學(xué)習(xí)領(lǐng)域，模型的參數(shù)、數(shù)據(jù)集的規(guī)模等等動(dòng)輒就是以億為單位，甚至更大，因此當(dāng)模型訓(xùn)練成功之時(shí)，放一首張靚穎的「終于等到你」作為背景音樂(lè)實(shí)在是太應(yīng)景了。
那如果現(xiàn)在向你推薦一款神器，可以實(shí)現(xiàn)訓(xùn)練速度翻倍，訪存效率翻倍，你心動(dòng)嗎？心動(dòng)不如行動(dòng)（這可不是電視直銷(xiāo)，別著急換頻道），來(lái)和我一起看看這款神器——基于飛槳核心框架的自動(dòng)混合精度（Automatic Mixed Precision) 技術(shù)，簡(jiǎn)稱飛槳 AMP 技術(shù)。
飛槳 AMP 技術(shù)僅僅通過(guò)一行代碼即可幫助用戶簡(jiǎn)便快速的將單精度訓(xùn)練的模型修改為自動(dòng)混合精度訓(xùn)練。同時(shí)通過(guò)黑白名單和動(dòng)態(tài) Loss Scaling 來(lái)保證訓(xùn)練的穩(wěn)定性，避免出現(xiàn) INF 或者 NAN 問(wèn)題。飛槳 AMP 可以充分發(fā)揮新一代 NVIDIA GPU 中 Tensor Core 的計(jì)算性能優(yōu)勢(shì)，ResNet50、Transformer 等模型的訓(xùn)練速度與單精度訓(xùn)練相比可以提升到 1.5～2.9 倍。
那么它是怎么實(shí)現(xiàn)的呢？我們先從什么是自動(dòng)混合精度技術(shù)講起。
什么是自動(dòng)混合精度技術(shù)
顧名思義，自動(dòng)混合精度是一種自動(dòng)將半精度和單精度混合使用，從而加速模型訓(xùn)練的技術(shù)。其中單精度（Float Precision32，F(xiàn)P32）好理解，是計(jì)算機(jī)常用的一種數(shù)據(jù)類型。那么半精度是什么呢？如圖 1 所示，半精度（Float Precision16，F(xiàn)P16）是一種相對(duì)較新的浮點(diǎn)類型，在計(jì)算機(jī)中使用 2 字節(jié)（16 位）存儲(chǔ)，在 IEEE 754-2008 中，它被稱作 binary16。與計(jì)算中常用的單精度和雙精度類型相比，F(xiàn)loat16 更適于在精度要求不高的場(chǎng)景中使用。
圖 1 半精度和單精度數(shù)據(jù)示意圖
不言而喻，在深度學(xué)習(xí)領(lǐng)域，如果使用 Float16 代替 Float32 來(lái)存儲(chǔ)數(shù)據(jù)，那么開(kāi)發(fā)者就可以訓(xùn)練更大更復(fù)雜的模型，使用更大的 batch size。因此對(duì)于那些恨不得挖掘出 GPU 里每一個(gè)晶體管全部潛力的科學(xué)家們?cè)趺茨芊胚^(guò)它呢？同時(shí)由于 NVIDIA 推出了具備 Tensor Core 技術(shù)的 Volta 及 Turing 架構(gòu) GPU，使半精度計(jì)算趨向成熟。在相同的 GPU 硬件上，Tensor Core 的半精度計(jì)算吞吐量是單精度的 8 倍。
但顯而易見(jiàn)，使用 Float16 肯定會(huì)同時(shí)帶來(lái)計(jì)算精度上的損失。但對(duì)深度學(xué)習(xí)訓(xùn)練而言，并不是所有計(jì)算都要求很高的精度，一些局部的精度損失對(duì)最終訓(xùn)練效果影響很微弱，僅需要某些特殊步驟保留 Float32 的計(jì)算精度即可。因此混合精度計(jì)算的需求應(yīng)運(yùn)而生。我們可以將訓(xùn)練過(guò)程中一些對(duì)精度損失不敏感且能使用 Tensor Core 進(jìn)行加速的運(yùn)算使用半精度處理，最大限度的提升訪存和計(jì)算效率。
但是對(duì)每個(gè)具體模型，人工去設(shè)計(jì)和嘗試精度混合的方法，是非常繁瑣的，我們迫切需要一種更簡(jiǎn)潔的方式，高效地實(shí)現(xiàn)混合精度的訓(xùn)練。AMP，顧名思義，就是讓混合精度訓(xùn)練自動(dòng)化，因此使用簡(jiǎn)單是它的重要特色。具體咋用，咱們往下看！
AMP 的使用方法
下面以 MNIST 為例介紹如何使用飛槳 AMP 技術(shù)。MNIST 網(wǎng)絡(luò)定義的代碼如下所示。其中 conv2d、batch_norm（bn）和 pool2d 的數(shù)據(jù)布局需要提前設(shè)置為'NHWC'，這樣有利于加速混合精度訓(xùn)練，并且 conv2d 的輸出通道數(shù)需要設(shè)置為 4 的倍數(shù)，以便使用 Tensor Core 技術(shù)加速。
import paddle.fluid as fluiddef MNIST(data, class_dim): conv1 = fluid.layers.conv2d(data, 16, 5, 1, act=None, data_format='NHWC') bn1 = fluid.layers.batch_norm(conv1, act='relu', data_layout='NHWC') pool1 = fluid.layers.pool2d(bn1, 2, 'max', 2, data_format='NHWC') conv2 = fluid.layers.conv2d(pool1, 64, 5, 1, act=None, data_format='NHWC') bn2 = fluid.layers.batch_norm(conv2, act='relu', data_layout='NHWC') pool2 = fluid.layers.pool2d(bn2, 2, 'max', 2, data_format='NHWC') fc1 = fluid.layers.fc(pool2, size=50, act='relu') fc2 = fluid.layers.fc(fc1, size=class_dim, act='softmax') return fc2
為了訓(xùn)練 MNIST 網(wǎng)絡(luò)，還需要定義損失函數(shù)來(lái)更新權(quán)重參數(shù)，此處使用的優(yōu)化損失函數(shù)是 SGDOptimizer。為了簡(jiǎn)化說(shuō)明，這里省略了迭代訓(xùn)練的相關(guān)代碼，僅體現(xiàn)損失函數(shù)及優(yōu)化器定義相關(guān)的內(nèi)容。
import paddle.fluid as fluidimport numpy as npdata = fluid.layers.data( name='image', shape=[None, 28, 28, 1], dtype='float32')label = fluid.layers.data(name='label', shape=[None, 1], dtype='int64')out = MNIST(data, class_dim=10)loss = fluid.layers.cross_entropy(input=out, label=label)avg_loss = fluid.layers.mean(loss)sgd = fluid.optimizer.SGDOptimizer(learning_rate=1e-3)sgd.minimize(avg_loss)
那么如何將上面的示例改造成使用 AMP 訓(xùn)練的方式呢？用戶僅需要使用飛槳提供的 AMP 函數(shù) fluid.contrib.mixed_precision.decorate 將原來(lái)的優(yōu)化器 SGDOptimizer 進(jìn)行封裝，然后使用封裝后的優(yōu)化器（mp_sgd）更新參數(shù)梯度，代碼如下所示：
sgd = fluid.optimizer.SGDOptimizer(learning_rate=1e-3)mp_sgd = fluid.contrib.mixed_precision.decorator.decorate(sgd)mp_sgd.minimize(avg_loss)
如上即為最簡(jiǎn)單的飛槳 AMP 功能使用方法。
但是大家可能有些疑問(wèn)，模型是如何感知哪些算子（Op）需要被轉(zhuǎn)換呢？是不是還需要手工指定呢？算子那么多，我怎么知道哪個(gè)算子可以被轉(zhuǎn)換呢？別著急，飛槳已經(jīng)幫你定制好了，這也是這門(mén)技術(shù)被稱為「自動(dòng)」的原因之一，且請(qǐng)往下看！
黑白名單功能
為了讓開(kāi)發(fā)者可以方便快捷的使用混合精度計(jì)算，飛槳的工程師們使用了大量模型在不同應(yīng)用場(chǎng)景中反復(fù)驗(yàn)證，然后根據(jù)半精度數(shù)據(jù)類型計(jì)算的穩(wěn)定性和加速效果，梳理出一系列適合轉(zhuǎn)換為半精度計(jì)算的算子，并將這些算子定義到了一份白名單文件中。同時(shí)對(duì)于一些經(jīng)過(guò)驗(yàn)證發(fā)現(xiàn)不適合轉(zhuǎn)換的算子，也就是使用半精度計(jì)算會(huì)導(dǎo)致數(shù)值不精確的算子將被記錄到黑名單文件中。此外一些對(duì)半精度計(jì)算沒(méi)有多少影響的算子歸類于灰名單。在使用 AMP 訓(xùn)練過(guò)程中，系統(tǒng)會(huì)自動(dòng)讀取黑白名單，從而感知到哪些算子需要被轉(zhuǎn)換為半精度計(jì)算。
對(duì)于某些特殊場(chǎng)景，如果開(kāi)發(fā)者希望使用自定義的黑白名單，則可以使用 AutoMixedPrecisionLists 類設(shè)置，代碼示例如下所示。
sgd = SGDOptimizer(learning_rate=1e-3)# 指定自定義的黑白名單，其中 list1 和 list2 為包含有算子名稱的列表amp_list = AutoMixedPrecisionLists(custom_white_list=list1,custom_black_list=list2)mp_sgd = fluid.contrib.mixed_precision.decorator.decorate(sgd, amp_list)mp_sgd.minimize(avg_loss)
那么自動(dòng)混合精度技術(shù)被稱為「自動(dòng)」的原因之二呢？那就是下面的自動(dòng)調(diào)整 Loss Scaling 功能。
自動(dòng)調(diào)整 Loss Scaling
AMP 技術(shù)在提升訪存和計(jì)算效率的同時(shí)，伴隨的副作用也是很明顯的。那就是由于半精度數(shù)據(jù)類型的精度范圍與轉(zhuǎn)換前的單精度相比過(guò)窄，導(dǎo)致容易產(chǎn)生 INF 和 NAN 問(wèn)題。為了避免此類問(wèn)題，AMP 技術(shù)實(shí)現(xiàn)了自動(dòng)調(diào)整 Loss Scaling 功能，即在 AMP 訓(xùn)練過(guò)程中，為了避免精度下溢，每訓(xùn)練一定數(shù)量批次的數(shù)據(jù)，就將 Loss 放大指定倍數(shù)。如果 Loss 在放大過(guò)程中發(fā)生上溢，則可以再縮小一定倍數(shù)，確保整個(gè)訓(xùn)練過(guò)程中，梯度可以正常收斂。
fluid.contrib.mixed_precision.decorate 函數(shù)攜帶了自動(dòng)調(diào)整 Loss Scaling 功能相關(guān)的參數(shù)，這些參數(shù)都帶有默認(rèn)值，如下面代碼所示。這些默認(rèn)值都是經(jīng)過(guò)飛槳工程師多次驗(yàn)證后定義的。通常情況下，用戶可以直接使用，無(wú)需重新設(shè)置。
sgd = SGDOptimizer(learning_rate=1e-3)mp_sgd = fluid.contrib.mixed_precision.decorator.decorate(sgd, init_loss_scaling=2**15, incr_every_n_steps=2000, use_dynamic_loss_scaling=True)mp_sgd.minimize(avg_loss)
多卡 GPU 訓(xùn)練的優(yōu)化
在新發(fā)布的飛槳核心框架 1.7 版本上，AMP 技術(shù)深度優(yōu)化了多卡 GPU 訓(xùn)練。如圖 2 所示，在優(yōu)化之前的參數(shù)梯度更新過(guò)程中，梯度計(jì)算時(shí)雖然使用的是半精度數(shù)據(jù)類型，但是不同 GPU 卡之間的梯度傳輸數(shù)據(jù)類型仍為單精度。
圖 2 1.7 版本之前的參數(shù)梯度更新過(guò)程示意圖
為了降低 GPU 多卡之間的梯度傳輸帶寬，我們將梯度傳輸這個(gè)過(guò)程提到 Cast 操作之前，而每個(gè) GPU 卡在得到對(duì)應(yīng)的半精度梯度后再執(zhí)行 Cast 操作，將其轉(zhuǎn)變?yōu)閱尉阮愋停鐖D 3 所示。這一優(yōu)化在訓(xùn)練網(wǎng)絡(luò)復(fù)雜度較大的模型時(shí)，對(duì)減少帶寬占用方面非常有效，如多卡訓(xùn)練 BERT-Large 模型。
圖 3 1.7 版本的參數(shù)梯度更新過(guò)程示意圖
訓(xùn)練性能對(duì)比（AMP VS FP32）
飛槳 AMP 技術(shù)在 ResNet50、Transformer 等模型上訓(xùn)練速度相對(duì)于 FP32 訓(xùn)練來(lái)說(shuō)有非常大的優(yōu)勢(shì)，下面以 ResNet50 模型為例，從下圖中可以看出，ResNet50 的 AMP 訓(xùn)練相對(duì)與 FP32 訓(xùn)練，單卡加速比可達(dá) 2.9 倍，八卡加速比可達(dá) 2.8 倍。
