用上Pytorch Lightning的這六招，深度學(xué)習(xí)pipeline提速10倍！

十三 2021-10-02 12:37:51 來源：量子位

金磊 發(fā)自 凹非寺

量子位 報道 | 公眾號 QbitAI

面對數(shù)以億計的圖片數(shù)據(jù)，到底該用什么樣的方法才能快速搞實驗？

這樣的問題，或許在做機器學(xué)習(xí)研究的你，也會經(jīng)常遇到。

而就在最近，一個國外小哥就提出了一種建議：

在Pytorch lightning基礎(chǔ)上，讓深度學(xué)習(xí)pipeline速度提升10倍！

用他自己的話來說就是——“爬樓時像給了你一個電梯”。

這般“酸爽”，到底是如何做到的呢？

優(yōu)化機器學(xué)習(xí)pipeline，很重要

無論你是身處學(xué)術(shù)界還是工業(yè)界，時間和資源等各種因素，往往會成為你在搞實驗的枷鎖。

尤其是隨著數(shù)據(jù)集規(guī)模和機器學(xué)習(xí)模型，變得越發(fā)龐大和復(fù)雜，讓實驗變得既費時又耗力。

提速這件事，就變得至關(guān)重要。

例如在2012年的時候，訓(xùn)練一個AlexNet，要花上5到6天的時間。

而現(xiàn)如今，只需要短短幾分鐘就可以在更大的數(shù)據(jù)集上訓(xùn)練更大的圖像模型。

這位小哥認為，從某種角度上來說，這是得益于各種各樣的“利器”的出現(xiàn)。

例如Pytorch Lingtning，就是其中一種。

于是，他便“死磕”pipeline，總結(jié)了六種“閃電加速”實驗周期的方法。

并行數(shù)據(jù)加載

數(shù)據(jù)加載和增強（augmentation）往往被認為是訓(xùn)練pipeline時的瓶頸之一。

一個典型的數(shù)據(jù)pipeline包含以下步驟：

• 從磁盤加載數(shù)據(jù)
• 在運行過程中創(chuàng)建隨機增強
• 將每個樣本分批整理

在這個過程中，倒是可以用多個CPU進程并行加載數(shù)據(jù)來優(yōu)化。

但與此同時，還可以通過下面的操作來加速這一過程：

1、將DataLoader中的num_workers參數(shù)設(shè)置為CPU的數(shù)量。

2、當與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”。

而小哥采用的方法是后者，因為他認為這是一種更可擴展的方法。

但在PyTorch（以及其他平臺）中修改訓(xùn)練pipeline并非易事。

必須考慮以分布式方式加載數(shù)據(jù)以及權(quán)重、梯度和指標的同步等問題。

不過，有了PyTorch Lightning，就可以非常容易地在多個GPU上訓(xùn)練PyTorch模型，還是幾乎不需要修改代碼的那種！

混合精度

在默認情況下，輸入張量以及模型權(quán)重是以單精度（float32）定義的。

然而，某些數(shù)學(xué)運算可以用半精度（float16）進行。

這樣一來，就可以顯著提升速度，并降低了模型的內(nèi)存帶寬，還不會犧牲模型的性能。

通過在PyTorch Lightning中設(shè)置混合精度標志（flag），它會在可能的情況下自動使用半精度，而在其他地方保留單精度。

通過最小的代碼修改，模型訓(xùn)練的速度可以提升1.5至2倍。

早停法

當我們訓(xùn)練深度學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)的時候，通常希望能獲得最好的泛化性能。

但是所有的標準深度學(xué)習(xí)神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)，比如全連接多層感知機都很容易過擬合。

當網(wǎng)絡(luò)在訓(xùn)練集上表現(xiàn)越來越好，錯誤率越來越低的時候，實際上在某一刻，它在測試集的表現(xiàn)已經(jīng)開始變差。

因此，早停法?（Early Stopping）便在訓(xùn)練過程中加入了進來。

具體來說，就是當驗證損失在預(yù)設(shè)的評估次數(shù)（在小哥的例子中是10次評估）后停止訓(xùn)練。

這樣一來，不僅防止了過擬合的現(xiàn)象，而且還可以在幾十個 epoch內(nèi)找到最佳模型。

Sharded Training

Sharded Training是基于微軟的ZeRO研究和DeepSpeed庫。

它顯著的效果，就是讓訓(xùn)練大模型變得可擴展和容易。

否則，這些模型就不適合在單個GPU上使用了。

而在Pytorch Lightning的1.2版本中，便加入了對Shared Training的支持。

雖然在小哥的實驗過程中，并沒有看到訓(xùn)練時間或內(nèi)存占用方面有任何改善。

但他認為，這種方法在其它實驗中可能會提供幫助，尤其是在不使用單一GPU的大模型方面。

模型評估和推理中的優(yōu)化

在模型評估和推理期間，梯度不需要用于模型的前向傳遞。

因此，可以將評估代碼包裹在一個torch.no_grad上下文管理器中。

這可以防止在前向傳遞過程中的存儲梯度，從而減少內(nèi)存占用。

如此一來，就可以將更大的batch送入模型，讓評估和推理變得更快。

效果如何？

介紹了這么多，你肯定想知道上述這些方法，具體起到了怎樣的作用。

小哥為此做了一張表格，詳解了方法的加速效果。

那么這些方法，是否對在做機器學(xué)習(xí)實驗的你有所幫助呢？

快去試試吧~

參考鏈接：

https://devblog.pytorchlightning.ai/how-we-used-pytorch-lightning-to-make-our-deep-learning-pipeline-10x-faster-731bd7ad318a