用上Pytorch Lightning的這六招,深度學(xué)習(xí)pipeline提速10倍
本文經(jīng)AI新媒體量子位(公眾號ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
面對數(shù)以億計的圖片數(shù)據(jù),到底該用什么樣的方法才能快速搞實驗?
這樣的問題,或許在做機(jī)器學(xué)習(xí)研究的你,也會經(jīng)常遇到。
而就在最近,一個國外小哥就提出了一種建議:
在Pytorch lightning基礎(chǔ)上,讓深度學(xué)習(xí)pipeline速度提升10倍!
用他自己的話來說就是——“爬樓時像給了你一個電梯”。
這般“酸爽”,到底是如何做到的呢?

優(yōu)化機(jī)器學(xué)習(xí)pipeline,很重要
無論你是身處學(xué)術(shù)界還是工業(yè)界,時間和資源等各種因素,往往會成為你在搞實驗的枷鎖。
尤其是隨著數(shù)據(jù)集規(guī)模和機(jī)器學(xué)習(xí)模型,變得越發(fā)龐大和復(fù)雜,讓實驗變得既費時又耗力。

提速這件事,就變得至關(guān)重要。
例如在2012年的時候,訓(xùn)練一個AlexNet,要花上5到6天的時間。
而現(xiàn)如今,只需要短短幾分鐘就可以在更大的數(shù)據(jù)集上訓(xùn)練更大的圖像模型。
這位小哥認(rèn)為,從某種角度上來說,這是得益于各種各樣的“利器”的出現(xiàn)。
例如Pytorch Lingtning,就是其中一種。
于是,他便“死磕”pipeline,總結(jié)了六種“閃電加速”實驗周期的方法。
并行數(shù)據(jù)加載
數(shù)據(jù)加載和增強(qiáng)(augmentation)往往被認(rèn)為是訓(xùn)練pipeline時的瓶頸之一。
一個典型的數(shù)據(jù)pipeline包含以下步驟:
- 從磁盤加載數(shù)據(jù)
- 在運行過程中創(chuàng)建隨機(jī)增強(qiáng)
- 將每個樣本分批整理
在這個過程中,倒是可以用多個CPU進(jìn)程并行加載數(shù)據(jù)來優(yōu)化。
但與此同時,還可以通過下面的操作來加速這一過程:
1、將DataLoader中的num_workers參數(shù)設(shè)置為CPU的數(shù)量。
2、當(dāng)與GPU一起工作時,將DataLoader中的pin_memory參數(shù)設(shè)置為True。這可以將數(shù)據(jù)分配到頁鎖定的內(nèi)存中,從而加快數(shù)據(jù)傳輸?shù)紾PU的速度。
使用分布式數(shù)據(jù)并行的多GPU訓(xùn)練

與CPU相比,GPU已經(jīng)大大加速了訓(xùn)練和推理時間。
但有沒有比一個GPU更好的方法?或許答案就是:
多個GPU!
在PyTorch中,有幾種范式可以用多個GPU訓(xùn)練你的模型。
兩個比較常見的范式是 “DataParallel ”和 “DistributedDataParallel”。
而小哥采用的方法是后者,因為他認(rèn)為這是一種更可擴(kuò)展的方法。
但在PyTorch(以及其他平臺)中修改訓(xùn)練pipeline并非易事。
必須考慮以分布式方式加載數(shù)據(jù)以及權(quán)重、梯度和指標(biāo)的同步等問題。
不過,有了PyTorch Lightning,就可以非常容易地在多個GPU上訓(xùn)練PyTorch模型,還是幾乎不需要修改代碼的那種!

混合精度
在默認(rèn)情況下,輸入張量以及模型權(quán)重是以單精度(float32)定義的。
然而,某些數(shù)學(xué)運算可以用半精度(float16)進(jìn)行。
這樣一來,就可以顯著提升速度,并降低了模型的內(nèi)存帶寬,還不會犧牲模型的性能。
通過在PyTorch Lightning中設(shè)置混合精度標(biāo)志(flag),它會在可能的情況下自動使用半精度,而在其他地方保留單精度。
通過最小的代碼修改,模型訓(xùn)練的速度可以提升1.5至2倍。

早停法
當(dāng)我們訓(xùn)練深度學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)的時候,通常希望能獲得最好的泛化性能。
但是所有的標(biāo)準(zhǔn)深度學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu),比如全連接多層感知機(jī)都很容易過擬合。
當(dāng)網(wǎng)絡(luò)在訓(xùn)練集上表現(xiàn)越來越好,錯誤率越來越低的時候,實際上在某一刻,它在測試集的表現(xiàn)已經(jīng)開始變差。
因此,早停法 (Early Stopping)便在訓(xùn)練過程中加入了進(jìn)來。
具體來說,就是當(dāng)驗證損失在預(yù)設(shè)的評估次數(shù)(在小哥的例子中是10次評估)后停止訓(xùn)練。
這樣一來,不僅防止了過擬合的現(xiàn)象,而且還可以在幾十個 epoch內(nèi)找到最佳模型。

Sharded Training
Sharded Training是基于微軟的ZeRO研究和DeepSpeed庫。
它顯著的效果,就是讓訓(xùn)練大模型變得可擴(kuò)展和容易。
否則,這些模型就不適合在單個GPU上使用了。
而在Pytorch Lightning的1.2版本中,便加入了對Shared Training的支持。
雖然在小哥的實驗過程中,并沒有看到訓(xùn)練時間或內(nèi)存占用方面有任何改善。
但他認(rèn)為,這種方法在其它實驗中可能會提供幫助,尤其是在不使用單一GPU的大模型方面。

模型評估和推理中的優(yōu)化
在模型評估和推理期間,梯度不需要用于模型的前向傳遞。
因此,可以將評估代碼包裹在一個torch.no_grad上下文管理器中。
這可以防止在前向傳遞過程中的存儲梯度,從而減少內(nèi)存占用。
如此一來,就可以將更大的batch送入模型,讓評估和推理變得更快。
效果如何?
介紹了這么多,你肯定想知道上述這些方法,具體起到了怎樣的作用。
小哥為此做了一張表格,詳解了方法的加速效果。

那么這些方法,是否對在做機(jī)器學(xué)習(xí)實驗的你有所幫助呢?
快去試試吧~






























