為了記錄這一過(guò)程，我們計(jì)算了網(wǎng)絡(luò)從零開始訓(xùn)練到94%的精確度所需的時(shí)間。這一基準(zhǔn)來(lái)自最近的DAWNBench競(jìng)賽。在競(jìng)賽結(jié)束后，單個(gè)GPU上的最好成績(jī)是341秒，八個(gè)GPU上最好成績(jī)是174秒。

Baseline

在這部分中，我們復(fù)制了一個(gè)基線，在6分鐘內(nèi)訓(xùn)練CIFAR10，之后稍稍加速。我們發(fā)現(xiàn)，在GPU的FLOPs計(jì)算完之前，仍有很大的提升空間。

過(guò)去幾個(gè)月，我一直在研究如何能更快度訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)。這個(gè)想法是從今年年初萌生的，當(dāng)時(shí)我正和Myrtle的SamDavis進(jìn)行一個(gè)項(xiàng)目。我們將用于自動(dòng)語(yǔ)音識(shí)別的大型循環(huán)神經(jīng)網(wǎng)絡(luò)壓縮后，部署到FPGAs上，重新訓(xùn)練模型。來(lái)自Mozilla的基線在16個(gè)GPU上訓(xùn)練了一個(gè)星期。后來(lái)，經(jīng)過(guò)Sam的努力，我們?cè)谟ミ_(dá)的VoltaGPUs上進(jìn)行混淆精度訓(xùn)練，得以將訓(xùn)練時(shí)間縮短了100倍，迭代時(shí)間在單個(gè)GPU上只需要不到一天的時(shí)間。

這一結(jié)果讓我思考還有什么可以實(shí)現(xiàn)加速？幾乎與此同時(shí)，斯坦福大學(xué)的研究人員們開啟了DAWNBench挑戰(zhàn)賽，比較多個(gè)深度學(xué)習(xí)基線上的訓(xùn)練速度。最受人關(guān)注的就是訓(xùn)練圖像分類模型在CIFAR10上達(dá)到94%的測(cè)試精確度，在ImageNet上達(dá)到93%、top5的成績(jī)。圖像分類是深度學(xué)習(xí)研究的熱門領(lǐng)域，但是訓(xùn)練速度仍需要數(shù)小時(shí)。

到了四月份，挑戰(zhàn)賽接近尾聲，CIFAR10上最快的單個(gè)GPU訓(xùn)練速度來(lái)自fast.ai的一名學(xué)生BenJohnson，他在不到6分鐘（341秒）的時(shí)間里訓(xùn)練出了94%的精確度。這一創(chuàng)新主要是混淆精度的訓(xùn)練，他選擇了一個(gè)較小的網(wǎng)絡(luò)，有足夠的能力處理任務(wù)并且可以用更高的學(xué)習(xí)速率加速隨機(jī)梯度下降。

這時(shí)我們不禁提出一個(gè)問(wèn)題：這種341秒訓(xùn)練出來(lái)的94%測(cè)試精度，在CIFAR10上的表現(xiàn)怎么樣？該網(wǎng)絡(luò)的架構(gòu)是一個(gè)18層的殘差網(wǎng)絡(luò)，如下所示。在這個(gè)案例中，圖層的數(shù)量表示卷積（紫色）和完全連接層（藍(lán)色）的序列深度：

網(wǎng)絡(luò)通過(guò)隨機(jī)梯度下降訓(xùn)練了35個(gè)epoch，學(xué)習(xí)速率圖如下：

現(xiàn)在我們假設(shè)在一個(gè)英偉達(dá)VoltaV100GPU上用100%的計(jì)算力，訓(xùn)練將需要多長(zhǎng)時(shí)間。網(wǎng)絡(luò)在一張32×32×3的CIFAR10圖像上進(jìn)行前向和后向傳遞時(shí)需要大約2.8×109FLOPs。假設(shè)參數(shù)更新不耗費(fèi)計(jì)算力，那么在50000張圖像訓(xùn)練35個(gè)epoch應(yīng)該會(huì)在5×1015FLOPs以內(nèi)完成。

TeslaV100有640個(gè)TensorCores，能支持125TeraFLOPS的深度學(xué)習(xí)性能。

假設(shè)我們能發(fā)揮100%的計(jì)算力，那么訓(xùn)練會(huì)在40秒內(nèi)完成，這么看來(lái)341秒的成績(jī)還有很大的提升空間。

有了40秒這個(gè)目標(biāo)，我們就開始了自己的訓(xùn)練。首先是用上方的殘差網(wǎng)絡(luò)重新復(fù)現(xiàn)基線CIFAR10的結(jié)果。我用PyTorch創(chuàng)建了一個(gè)網(wǎng)絡(luò)，重新復(fù)制了學(xué)習(xí)速率和超參數(shù)。在AWSp3.2的圖像上用單個(gè)V100GPU訓(xùn)練，3/5的運(yùn)行結(jié)果在356秒內(nèi)達(dá)到了94%的精確度。

基線建好后，下一步是尋找可以立即使用的簡(jiǎn)單改進(jìn)方法。首先我們觀察到：網(wǎng)絡(luò)開頭是由黃色和紅色的兩個(gè)連續(xù)norm-ReLU組成的，在紫色卷積之后，我們刪去重復(fù)部分，同樣在epoch15也發(fā)生了這樣的情況。進(jìn)行調(diào)整后，網(wǎng)絡(luò)架構(gòu)變得更簡(jiǎn)單，4/5的運(yùn)行結(jié)果在323秒內(nèi)達(dá)到了94%的精確度！刷新了記錄！

另外我們還觀察到，圖像處理過(guò)程中的一些步驟（填充、標(biāo)準(zhǔn)化、位移等等）每經(jīng)過(guò)訓(xùn)練集一次就要重新處理一遍，會(huì)浪費(fèi)很多時(shí)間。雖然提前預(yù)處理可以用多個(gè)CPU處理器減輕這一結(jié)果，但是PyTorch的數(shù)據(jù)下載器會(huì)從每次數(shù)據(jù)迭代中開始新一次的處理。這一配置時(shí)間是很短的，尤其在CIFAR10這樣的小數(shù)據(jù)集上。只要在訓(xùn)練前做了準(zhǔn)備，減少預(yù)處理壓力，就能減少處理次數(shù)。遇到更復(fù)雜的任務(wù)，需要更多預(yù)處理步驟或多個(gè)GPU時(shí)，就會(huì)在每個(gè)epoch之間保持?jǐn)?shù)據(jù)下載器的處理。溢出了重復(fù)工作、減少了數(shù)據(jù)下載器后，訓(xùn)練時(shí)間達(dá)到了308秒。

繼續(xù)研究后我們發(fā)現(xiàn)，大部分預(yù)處理時(shí)間都花在了召集隨機(jī)數(shù)字生成器，選擇數(shù)據(jù)增強(qiáng)而不是為它們本身增強(qiáng)。在完全訓(xùn)練時(shí)期，我們對(duì)隨機(jī)數(shù)字生成器執(zhí)行了幾百萬(wàn)個(gè)單獨(dú)命令，把它們結(jié)合在一個(gè)較小的命令中，每個(gè)epoch可以省去7秒訓(xùn)練時(shí)間。