26秒訓(xùn)練ResNet，用這些技巧一步步壓縮時(shí)間，Jeff Dean都稱贊：干得漂亮

魚羊 乾明 發(fā)自 凹非寺

量子位 報(bào)道 | 公眾號(hào) QbitAI

更快的訓(xùn)練速度，更少的算力消耗，對(duì)于煉丹師們而言，這無(wú)疑是飛一般的體驗(yàn)。

現(xiàn)在，谷歌AI掌門人Jeff Dean轉(zhuǎn)發(fā)推薦了一個(gè)訓(xùn)練ResNet的奇技淫巧大禮包，跟著它一步一步實(shí)施，訓(xùn)練9層ResNet時(shí)，不僅不需要增加GPU的數(shù)量，甚至只需要1/8的GPU，就能讓訓(xùn)練速度加快到原來(lái)的2.5倍，模型在CIFAR10上還能達(dá)到94%的準(zhǔn)確率。

甚至只需要26秒，就能訓(xùn)練好一個(gè)模型。

這一”大禮包“由myrtle.ai開源。Jeff Dean稱贊道：

真是減少CIFAR10圖像模型訓(xùn)練時(shí)間的好文章。其中許多技巧都可能適用于其他不同類型的模型。干得漂亮。

奇技淫巧大禮包

myrtle.ai的奇技淫巧，都可以在Colab上測(cè)試，只需配置一個(gè)V100。（鏈接見文末）

話不多說(shuō)，快來(lái)看看都是哪些神奇的技巧，造就飛一般的感覺。

基線方法消耗的總時(shí)間是75s。

GPU上的預(yù)處理（70s）

第一個(gè)技巧是，將數(shù)據(jù)傳輸?shù)紾PU，在GPU上進(jìn)行預(yù)處理，然后再傳回CPU進(jìn)行隨機(jī)數(shù)據(jù)擴(kuò)增和批處理。

將整個(gè)數(shù)據(jù)集（unit8格式）移至GPU僅需要40ms，幾乎可以忽略不計(jì)，而在GPU上完成預(yù)處理步驟花費(fèi)的時(shí)間則更少，大概只需要15ms。

這種做法下大部分的時(shí)間是消耗在數(shù)據(jù)集回傳CPU上，這需要將近500ms的時(shí)間。

往常的數(shù)據(jù)預(yù)處理方法會(huì)花費(fèi)掉3s以上的時(shí)間，相比之下，這樣的處理方式已經(jīng)把速度加快了不少。

不過(guò)，還可以更快。

那就是，不回傳CPU，直接在GPU上把數(shù)據(jù)擴(kuò)增這一步也做了。

當(dāng)然，蠻干是不行的。為了避免啟動(dòng)多個(gè)GPU內(nèi)核導(dǎo)致花銷變大，可以對(duì)樣本組應(yīng)用相同的擴(kuò)增，并通過(guò)預(yù)先對(duì)數(shù)據(jù)進(jìn)行混洗的方式來(lái)保持隨機(jī)性。

具體的做法是這樣的：

以8×8的cutout為例，即對(duì)CIFAR10圖像做8×8的隨機(jī)塊狀遮擋。

在32×32的圖像中有625個(gè)可能的8×8剪切區(qū)域，因此通過(guò)混洗數(shù)據(jù)集，將其分成625個(gè)組，每個(gè)組代表一個(gè)剪切區(qū)域，即可實(shí)現(xiàn)隨機(jī)擴(kuò)增。

事實(shí)上，選擇大小均勻的組，與為每個(gè)樣本進(jìn)行隨機(jī)選擇并不完全相同，但也很接近。不過(guò)為了進(jìn)一步優(yōu)化，如果應(yīng)用同一種擴(kuò)增的組的數(shù)量太大，可以對(duì)其設(shè)置一個(gè)合理的限制范圍。

如此一來(lái)，迭代24個(gè)epoch，并對(duì)其進(jìn)行隨機(jī)裁減、水平翻轉(zhuǎn)、cutout數(shù)據(jù)擴(kuò)增，以及數(shù)據(jù)混洗和批處理，只需要不到400ms。

還有一個(gè)好處是，CPU預(yù)處理隊(duì)列和GPU不用再相互賽跑，這樣就不必再擔(dān)心數(shù)據(jù)加載的問(wèn)題了。

需要注意的是，這樣操作的前提是數(shù)據(jù)集足夠小，可以在GPU內(nèi)存中作為一個(gè)整體進(jìn)行存儲(chǔ)和操作。但Nvidia DALI這樣的工業(yè)強(qiáng)度解決方案或許可以實(shí)現(xiàn)進(jìn)一步的突破。

移動(dòng)最大池化層（64s）

最大池化和ReLU的順序是可以交換的。經(jīng)典的卷積池化是這樣：

可以把它調(diào)整成這樣：

切換順序?qū)⑹?4個(gè)epoch的訓(xùn)練時(shí)間進(jìn)一步減少3秒，而網(wǎng)絡(luò)功能完全沒有變化。

甚至還可以進(jìn)一步把池化提前：

這能進(jìn)一步使訓(xùn)練時(shí)間縮短5秒，但會(huì)導(dǎo)致網(wǎng)絡(luò)的改變。不過(guò)實(shí)驗(yàn)表明，這樣做對(duì)測(cè)試精度的負(fù)面影響很小，基線是94.1%，池化前移后會(huì)準(zhǔn)確率會(huì)降到94.0%（50次運(yùn)算的平均值），僅有0.1%的精度損失。

標(biāo)簽平滑（59s）

標(biāo)簽平滑是提高分類問(wèn)題中神經(jīng)網(wǎng)絡(luò)訓(xùn)練速度和泛化的一個(gè)成熟技巧。

應(yīng)用標(biāo)簽平滑之后，測(cè)試精度能提高到94.2%（50次運(yùn)算的平均值），這樣一來(lái)，就可以通過(guò)減少epoch的數(shù)量來(lái)?yè)Q取訓(xùn)練速度的提升。

CELU激活（52s）

平滑的激活函數(shù)對(duì)于優(yōu)化過(guò)程也很有幫助。

因此可以選擇連續(xù)可微分指數(shù)線性單元或者CELU激活，來(lái)替代ReLU。

當(dāng)平滑參數(shù)α為0.075的時(shí)候，測(cè)試精度能達(dá)到94.3%，于是，與標(biāo)簽平滑相比，又可以進(jìn)一步減少epoch，速度也就從59s縮短到了52s。

Ghost批量歸一（46s）

批量歸一最合適的批量大小大概在32左右。

但在批量大小比較大的時(shí)候，比如512，降低其大小會(huì)嚴(yán)重影響訓(xùn)練時(shí)間。不過(guò)這一問(wèn)題可以通過(guò)對(duì)batch的子集分別進(jìn)行批量歸一來(lái)解決，這種方法稱為“ghost”批量歸一。

固定批量歸一規(guī)模（43s）

批量歸一規(guī)范了每個(gè)通道的均值和方差，但這取決于可學(xué)習(xí)的規(guī)模和偏差。

如果通道規(guī)模發(fā)生很大的變化，就可能會(huì)減少通道的有效數(shù)量而造成瓶頸。

從上面這兩張圖中可以看出，規(guī)模并沒有進(jìn)行太多學(xué)習(xí)，主要是在權(quán)重衰減的控制下進(jìn)化。嘗試將規(guī)模固定在1/4的恒定值（訓(xùn)練中點(diǎn)的平均值）。最后一層的可學(xué)習(xí)比例稍大，這可以通過(guò)調(diào)整網(wǎng)絡(luò)輸出的比例來(lái)進(jìn)行補(bǔ)償。

實(shí)際上，如果將CELU的α參數(shù)重新調(diào)整為補(bǔ)償因子4，批量歸一偏差的學(xué)習(xí)率和權(quán)重衰減分別為4^2和（1/4）^2，則批量歸一規(guī)模就為1。

需要說(shuō)明的是，如果不提高批量歸一偏差的學(xué)習(xí)率，最終的準(zhǔn)確率會(huì)非常低。

至此，在DAWNBench排行榜上，這一單GPU訓(xùn)練出來(lái)的ResNet9訓(xùn)練速度已經(jīng)能超越8個(gè)GPU訓(xùn)練出來(lái)的BaiduNet9了。

輸入塊（patch）白化（36s）

批量歸一可以很好地控制各個(gè)通道呃分布，但不能解決通道和像素之間協(xié)方差的問(wèn)題。所以，要引入“白化”版本來(lái)控制內(nèi)部層的協(xié)方差。

myrtle.ai提出了一種基于patch的方法，該方法與總圖像的尺寸無(wú)關(guān)，并且更符合卷積網(wǎng)絡(luò)的結(jié)構(gòu)。

將PCA白化應(yīng)用于3×3的輸入塊，作為具有固定（不可學(xué)習(xí)）權(quán)重的初始3×3卷積。這可以通過(guò)可學(xué)習(xí)的1×1卷積來(lái)實(shí)現(xiàn)。該層的27個(gè)輸入通道是原始的3×3×3輸入塊的變形，其協(xié)方差矩陣近于恒等式（identity），更易于優(yōu)化。

首先，繪制輸入數(shù)據(jù)的3×3塊的協(xié)方差矩陣的前導(dǎo)特征向量。括號(hào)中的數(shù)字是相應(yīng)特征值的平方根，以顯示沿這些方向的相對(duì)變化尺度，并繪制具有兩個(gè)符號(hào)的特征向量以說(shuō)明變化的方向。

可以看出，局部亮度的變化占據(jù)主導(dǎo)地位。

接下來(lái)，用固定的3×3白化卷積替換網(wǎng)絡(luò)的第一個(gè)3×3卷積，以均衡本征塊的比例，而后用可學(xué)習(xí)的1×1卷積，并查看其對(duì)訓(xùn)練的影響。

如果將最大學(xué)習(xí)率提高50%，并把剪切擴(kuò)增的總量從8×8降至5×5，以補(bǔ)償更高的學(xué)習(xí)率帶來(lái)的額外正則化，就可以在36s內(nèi)讓模型達(dá)到94.1的測(cè)試精度。

指數(shù)移動(dòng)平均(34s)

高學(xué)習(xí)率，是快速訓(xùn)練的必要條件，因?yàn)樗试S隨機(jī)梯度下降在有限的時(shí)間內(nèi)在參數(shù)空間中通過(guò)必要的距離。

另一方面，學(xué)習(xí)速率需要在訓(xùn)練結(jié)束時(shí)進(jìn)行退火，以便在參數(shù)空間中沿著更陡峭和更嘈雜的方向進(jìn)行優(yōu)化。

參數(shù)平均方法允許以更快的速度繼續(xù)訓(xùn)練，同時(shí)通過(guò)多次迭代進(jìn)行平均，可以沿著有噪聲或振蕩的方向接近最小值。

myrtle.ai發(fā)現(xiàn)，更頻繁的更新并不能改善情況，出于效率的原因，就每5個(gè)batch更新一次移動(dòng)平均。

他們說(shuō)，這需要選擇一個(gè)新的學(xué)習(xí)率計(jì)劃，在訓(xùn)練接近尾聲時(shí)提高學(xué)習(xí)率，并為移動(dòng)平均提供動(dòng)力。

對(duì)于學(xué)習(xí)速率，一個(gè)簡(jiǎn)單的選擇就是堅(jiān)持他們一直使用的分段線性時(shí)間表，在最后兩個(gè)epoch以一個(gè)較低的固定值表示，他們選擇的是一個(gè)0.99的動(dòng)量，這樣平均就可以在大約上一個(gè)時(shí)期的時(shí)間尺度上進(jìn)行。

測(cè)試精度提高到94.3%，也意味著可以進(jìn)一步削減epoch。

最后的結(jié)果是：13個(gè)epoch訓(xùn)練達(dá)到94.1%的測(cè)試精度，訓(xùn)練時(shí)間低于34秒，比這一系列開始的時(shí)候單GPU表現(xiàn)提高了10倍！

測(cè)試狀態(tài)增強(qiáng)(26秒)

如果你想讓你的網(wǎng)絡(luò)在輸入的水平翻轉(zhuǎn)下，也能以相同的方式對(duì)圖像進(jìn)行分類。

一種方法是向網(wǎng)絡(luò)提供大量的數(shù)據(jù)，通過(guò)保存左右翻轉(zhuǎn)的標(biāo)簽來(lái)增強(qiáng)數(shù)據(jù)，然后希望網(wǎng)絡(luò)最終通過(guò)廣泛的訓(xùn)練來(lái)學(xué)習(xí)不變性。

第二種方法，是同時(shí)呈現(xiàn)輸入圖像和水平翻轉(zhuǎn)后的圖像，并通過(guò)對(duì)兩個(gè)版本的網(wǎng)絡(luò)輸出進(jìn)行平均來(lái)達(dá)成一致，從而保證不變性。

這種非常明智的方法被稱為測(cè)試狀態(tài)增強(qiáng)（Test-time augmentation，TTA）。

在訓(xùn)練時(shí)，myrtle.ai仍然向網(wǎng)絡(luò)呈現(xiàn)每個(gè)圖像的單一版本——可能會(huì)進(jìn)行隨機(jī)翻轉(zhuǎn)來(lái)增強(qiáng)數(shù)據(jù)，以便在不同的訓(xùn)練階段呈現(xiàn)不同的版本。

他們采取的另一個(gè)做法是，在訓(xùn)練時(shí)使用與測(cè)試時(shí)相同的程序，并將每個(gè)圖像及其鏡像顯示出來(lái)。

在這種情況下，myrtle.ai可以通過(guò)將網(wǎng)絡(luò)分成兩個(gè)相同的分支來(lái)改變網(wǎng)絡(luò)，其中一個(gè)分支可以看到翻轉(zhuǎn)后的圖像，然后在最后合并。

通過(guò)這一視角，原始訓(xùn)練可以被看作是一個(gè)權(quán)重綁定的隨機(jī)訓(xùn)練過(guò)程，即兩個(gè)分支網(wǎng)絡(luò)，其中每個(gè)訓(xùn)練示例都有一個(gè)分支被”dropped-out“。

這種dropout-訓(xùn)練的觀點(diǎn)清楚地表明，任何試圖引入禁止測(cè)試狀態(tài)增強(qiáng)從基準(zhǔn)中刪除的規(guī)則都將充滿困難。

從這個(gè)角度看，myrtle.ai剛剛介紹了一個(gè)更大的網(wǎng)絡(luò)，他們有一個(gè)有效的隨機(jī)訓(xùn)練方法。

另一方面，如果myrtle.ai不限制準(zhǔn)備在測(cè)試時(shí)做的工作量，那么就會(huì)出現(xiàn)一些明顯的退化解決方案，其中訓(xùn)練所需的時(shí)間與存儲(chǔ)數(shù)據(jù)集所需的時(shí)間一樣少！

這些參數(shù)不僅與人工基準(zhǔn)測(cè)試相關(guān)，而且與最終的用例有關(guān)。在一些應(yīng)用程序中，對(duì)分類精度有所要求，在這種情況下，絕對(duì)應(yīng)該使用測(cè)試狀態(tài)增強(qiáng)。

在其他情況下，推斷時(shí)間也是一種約束，明智的做法是在這種約束下，將精度最大化。這可能也是一個(gè)訓(xùn)練基準(zhǔn)測(cè)試的好方法。

在目前的案例中，Kakao Brain團(tuán)隊(duì)?wèi)?yīng)用了這里描述的測(cè)試狀態(tài)增強(qiáng)的簡(jiǎn)單形式——在推斷時(shí)呈現(xiàn)圖像及其左右鏡像，從而使計(jì)算量加倍。

當(dāng)然，對(duì)于其他對(duì)稱(如平移對(duì)稱、亮度/顏色的變化等)來(lái)說(shuō)，測(cè)試狀態(tài)增強(qiáng)的應(yīng)用也會(huì)更廣泛，但這將付出更高的計(jì)算成本。

現(xiàn)在，由于目前是基于一個(gè)計(jì)算量較少的9層ResNet，因此包括測(cè)試狀態(tài)增強(qiáng)在內(nèi)的總推斷時(shí)間，可能比100多個(gè)層網(wǎng)絡(luò)中要短得多。

根據(jù)上面的討論，任何限制這種方法的合理規(guī)則，都應(yīng)該基于推斷時(shí)間約束，而不是實(shí)現(xiàn)的功能，因此從這個(gè)角度來(lái)看，myrtle.ai團(tuán)隊(duì)說(shuō)，應(yīng)該接受這種方法。

為了與當(dāng)前DAWNBench提交數(shù)值的一致，他們將限制自己使用水平翻轉(zhuǎn)測(cè)試狀態(tài)增強(qiáng)，因?yàn)檫@似乎是準(zhǔn)確性和推斷成本之間的最優(yōu)平衡點(diǎn)。

最后的結(jié)果是：在現(xiàn)有的網(wǎng)絡(luò)和13個(gè)epoch訓(xùn)練設(shè)置下，測(cè)試狀態(tài)增強(qiáng)的精度提高到94.6%。

如果移除了對(duì)剩余數(shù)據(jù)的增強(qiáng)操作，可以將訓(xùn)練減少到 10 個(gè) epoch，而且能在26 秒實(shí)現(xiàn)測(cè)試狀態(tài)增強(qiáng)精度——94.1%。

訓(xùn)練效果

單塊GPU就能有這么快的速度，效果又如何呢？myrtle.ai團(tuán)隊(duì)還用了ImageNet來(lái)驗(yàn)證。

從24個(gè)epoch到100個(gè)epoch，實(shí)驗(yàn)?zāi)Ｐ偷谋憩F(xiàn)始終優(yōu)于基線方法。

用上了奇技淫巧大禮包的9層ResNet其實(shí)去年11月就登上了DAWNBench CIFAR10排行榜的榜首，速度提高近2.5倍，而GPU從8個(gè)降到了1個(gè)。

DAWNBench是斯坦福大學(xué)提出的基準(zhǔn)，在這一排行榜中，準(zhǔn)確度只要達(dá)到94%即可。

現(xiàn)在，雖然這一登頂成績(jī)已經(jīng)滑落到了第六位，但myrtle.ai表示在使用了與榜首Kakao Brain相同的TTA方法之后，實(shí)驗(yàn)?zāi)Ｐ偷挠?xùn)練速度能降至26秒，超過(guò)Kakao Brain近10秒。

傳送門

博客地址：

How to Train Your ResNet 8: Bag of Tricks

Colab地址：

https://colab.research.google.com/github/davidcpage/cifar10-fast/blob/master/bag_of_tricks.ipynb#scrollTo=n___bs94Rvm2

GitHub地址：

https://github.com/davidcpage/cifar10-fast/blob/master/bag_of_tricks.ipynb

— 完 —