為了不讓GPU等CPU，谷歌提出“數(shù)據(jù)回波”榨干GPU空閑時(shí)間，訓(xùn)練速度提升3倍多

曉查 2020-05-14 13:32:59 來源：量子位

曉查 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI

因?yàn)橥ㄓ糜?jì)算芯片不能滿足神經(jīng)網(wǎng)絡(luò)運(yùn)算需求，越來越多的人轉(zhuǎn)而使用GPU和TPU這類專用硬件加速器，加快神經(jīng)網(wǎng)絡(luò)訓(xùn)練的速度。

但是，用了更快的GPU和TPU就一定能加速訓(xùn)練嗎？

訓(xùn)練流水線的所有操作并不都是在加速器上運(yùn)行。上游數(shù)據(jù)處理（如磁盤I/O和數(shù)據(jù)預(yù)處理）就不能在加速器上運(yùn)行。

隨著GPU等加速器越來越快，超過了CPU、磁盤處理數(shù)據(jù)的速度，上游就逐漸成了訓(xùn)練瓶頸。

在某些情況下，GPU上游的代碼花費(fèi)的時(shí)間甚至是GPU本身運(yùn)行時(shí)間的幾倍。上游沒做完，下游只能空等，浪費(fèi)了大量時(shí)間。

為此，Google AI團(tuán)隊(duì)，提出一種簡單的數(shù)據(jù)回波（Data Echoing）方法，可以解決這個(gè)問題。該方法最多能將訓(xùn)練速度加速3倍以上，且對(duì)模型的精度沒有明顯影響。

Jeff Dean也在Twitter上轉(zhuǎn)發(fā)點(diǎn)贊。

重復(fù)數(shù)據(jù)讓GPU不空等

很多情況下，上游花費(fèi)的時(shí)間比加速器更長，使用更快的GPU或TPU根本不會(huì)提高訓(xùn)練速度。如果投入大量的工程工作以及額外的計(jì)算資源，確實(shí)可以加快流水線的速度。

對(duì)于非常小的數(shù)據(jù)集，可以離線預(yù)先計(jì)算擴(kuò)增的數(shù)據(jù)集，并將整個(gè)預(yù)處理的數(shù)據(jù)集加載到內(nèi)存中。

但這種方法不適用于大多數(shù)機(jī)器學(xué)習(xí)訓(xùn)練場景，既耗時(shí)又分散了改善推理性能的主要目標(biāo)。

與其等待更多的數(shù)據(jù)，不如利用已有的數(shù)據(jù)來使加速器保持忙碌狀態(tài)。

在加速器空置50%情況下，預(yù)處理batch的第一個(gè)優(yōu)化步驟之后，我們可以重復(fù)利用該batch再進(jìn)行一次訓(xùn)練。

如果重復(fù)數(shù)據(jù)與新數(shù)據(jù)一樣有用，訓(xùn)練效率會(huì)提高一倍。

實(shí)際上，由于重復(fù)數(shù)據(jù)不如新數(shù)據(jù)有用，因此數(shù)據(jù)回波提供的加速要小一些，但和加速器處于空閑狀態(tài)相比，仍然可以提供明顯的加速。

通常有幾種方法可以在給定的神經(jīng)網(wǎng)絡(luò)訓(xùn)練管道中實(shí)現(xiàn)數(shù)據(jù)回波。

Google提出的技術(shù)，是將數(shù)據(jù)復(fù)制到訓(xùn)練管道中某個(gè)位置的隨機(jī)緩沖區(qū)中，無論在哪個(gè)階段產(chǎn)生瓶頸之后，都可以將緩存數(shù)據(jù)插入任意位置。

數(shù)據(jù)回波在樣本級(jí)別對(duì)數(shù)據(jù)進(jìn)行混洗，而batch回波則對(duì)重復(fù)批次的序列進(jìn)行混洗。另外還可以在數(shù)據(jù)擴(kuò)充之前插入緩沖區(qū)，以使重復(fù)數(shù)據(jù)的每個(gè)副本略有不同，因此不是簡單機(jī)械重復(fù)，而是更接近一個(gè)新樣本。

加速3倍多，精度無損失

那么數(shù)據(jù)回波到底多有用呢？

Google在五個(gè)神經(jīng)網(wǎng)絡(luò)訓(xùn)練管道上嘗試了數(shù)據(jù)回波，涵蓋了圖像分類、語言建模和目標(biāo)檢測3個(gè)不同的任務(wù)，并測量了達(dá)到特定性能目標(biāo)所需的新樣本數(shù)量。

Google發(fā)現(xiàn)，數(shù)據(jù)回波可以用更少的新樣本來達(dá)到目標(biāo)性能，這表明重復(fù)使用數(shù)據(jù)對(duì)于降低磁盤I/O很有用。在某些情況下，重復(fù)數(shù)據(jù)幾乎與新數(shù)據(jù)一樣有用。

在ImageNet的ResNet-50訓(xùn)練任務(wù)中，數(shù)據(jù)回波可以顯著提高訓(xùn)練速度，加速了3倍以上。

從云存儲(chǔ)中讀取一批訓(xùn)練數(shù)據(jù)所花的時(shí)間是使用每一批數(shù)據(jù)執(zhí)行訓(xùn)練步驟的6倍。因此數(shù)據(jù)最多可以重復(fù)5次。

我們將回波因子定義為重復(fù)每個(gè)數(shù)據(jù)項(xiàng)的次數(shù)，對(duì)于以上任務(wù)，回波因子最大是5。如果重復(fù)樣本與新樣本一樣有用，則應(yīng)該帶來6倍加速，而實(shí)際上只加速了3.25倍，但也相當(dāng)可觀。

可能有人會(huì)擔(dān)心重復(fù)使用數(shù)據(jù)會(huì)損害模型的最終性能，但實(shí)驗(yàn)發(fā)現(xiàn)，測試的任何工作負(fù)載，數(shù)據(jù)回波都不會(huì)降低最終模型的質(zhì)量。

隨著GPU和TPU性能的繼續(xù)提升，和通用處理器的差距會(huì)越來越大，Google期望數(shù)據(jù)回波和類似策略將成為神經(jīng)網(wǎng)絡(luò)培訓(xùn)工具包中越來越重要的一部分。

怎么樣，谷歌AI的這項(xiàng)最新研究，是不是也給你訓(xùn)練神經(jīng)網(wǎng)絡(luò)帶來一些新啟發(fā)？

論文地址：https://arxiv.org/abs/1907.05550

博客地址：https://ai.googleblog.com/2020/05/speeding-up-neural-network-training.html