Batch大小不一定是2的n次冪！ML資深學(xué)者最新結(jié)論
2的n次冪就一定能加速嗎？
羿閣 編譯整理
量子位 | 公眾號 QbitAI
Batch大小不一定是2的n次冪？
是否選擇2的n次冪在運行速度上竟然也相差無幾？
有沒有感覺常識被顛覆？
這是威斯康星大學(xué)麥迪遜分校助理教授Sebastian Raschka（以下簡稱R教授）的最新結(jié)論。
在神經(jīng)網(wǎng)絡(luò)訓(xùn)練中，2的n次冪作為Batch大小已經(jīng)成為一個標準慣例，即64、128、256、512、1024等。
一直有種說法，是這樣有助于提高訓(xùn)練效率。
但R教授做了一番研究之后，發(fā)現(xiàn)并非如此。
在介紹他的試驗方法之前，首先來回顧一下這個慣例究竟是怎么來的？
2的n次冪從何而來？
一個可能的答案是：因為CPU和GPU的內(nèi)存架構(gòu)都是由2的n次冪構(gòu)成的。
或者更準確地說，根據(jù)內(nèi)存對齊規(guī)則，cpu在讀取內(nèi)存時是一塊一塊進行讀取的，塊的大小可以是2，4，8，16（總之是2的倍數(shù)）。
因此，選取2的n次冪作為batch大小，主要是為了將一個或多個批次整齊地安裝在一個頁面上，以幫助GPU并行處理。
其次，矩陣乘法和GPU計算效率之間也存在一定的聯(lián)系。
假設(shè)我們在矩陣之間有以下矩陣乘法A和B：
當(dāng)A的行數(shù)等于B的列數(shù)的時候，兩個矩陣才能相乘。
其實就是矩陣A的第一行每個元素分別與B的第一列相乘再求和，得到C矩陣的第一個數(shù)，然后A矩陣的第一行再與B矩陣的第二列相乘，得到第二個數(shù)，然后是A矩陣的第二行與B矩陣的第一列……
因此，如上圖所示，我們擁有2×M×N×K個每秒浮點運算次數(shù)（FLOPS）。
現(xiàn)在，如果我們使用帶有Tensor Cores的GPU，例如V100時，當(dāng)矩陣尺寸（M，N以及K）與16字節(jié)的倍數(shù)對齊，在FP16混合精度訓(xùn)練中，8的倍數(shù)的運算效率最為理想。
因此，假設(shè)在理論上，batch大小為8倍數(shù)時，對于具有Tensor Cores和FP16混合精度訓(xùn)練的GPU最有效，那么讓我們調(diào)查一下這一說法在實踐中是否也成立。
不用2的n次冪也不影響速度
為了了解不同的batch數(shù)值對訓(xùn)練速度的影響，R教授在CIFAR-10上運行了一個簡單的基準測試訓(xùn)練——MobileNetV3（大）——圖像的大小為224×224，以便達到適當(dāng)?shù)腉PU利用率。
R教授用16位自動混合精度訓(xùn)練在V100卡上運行訓(xùn)練，該訓(xùn)練能更高效地使用GPU的Tensor Cores。
如果你想自己運行，該代碼可在此GitHub存儲庫中找到（鏈接附在文末）。
該測試共分為以下三部分：
小批量訓(xùn)練
從上圖可以看出，以樣本數(shù)量128為參考點，將樣本數(shù)量減少1（127）或增加1（129），的確會導(dǎo)致訓(xùn)練速度略慢，但這種差異幾乎可以忽略不計。
而將樣本數(shù)量減少28（100）會導(dǎo)致訓(xùn)練速度明顯放緩，這可能是因為模型現(xiàn)在需要處理的批次比以前更多（50,000/100=500與50,000/128= 390）。
同樣的原理，當(dāng)我們將樣本數(shù)量增加28（156）時，運行速度明顯變快了。
最大批量訓(xùn)練
鑒于MobileNetV3架構(gòu)和輸入映像大小，上一輪中樣本數(shù)量相對較小，因此GPU利用率約為70%。
為了調(diào)查GPU滿載時的訓(xùn)練速度，本輪把樣本數(shù)量增加到512，使GPU的計算利用率接近100%。
△由于GPU內(nèi)存限制，無法使用大于515的樣本數(shù)量
可以看出，跟上一輪結(jié)果一樣，不管樣本數(shù)量是否是2的n次冪，訓(xùn)練速度的差異幾乎可以忽略不計。
多GPU訓(xùn)練
基于前兩輪測試評估的都是單個GPU的訓(xùn)練性能，而如今多個GPU上的深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練更常見。為此，這輪進行的是多GPU培訓(xùn)。
正如我們看到的，2的n次冪（256）的運行速度并不比255差太多。
測試注意事項
在上述3個基準測試中，需要特別聲明的是：
所有基準測試的每個設(shè)置都只運行過一次，理想情況下當(dāng)然是重復(fù)運行次數(shù)越多越好，最好還能生成平均和標準偏差，但這并不會影響到上述結(jié)論。
此外，雖然R教授是在同一臺機器上運行的所有基準測試，但兩次運營之間沒有特意相隔很長時間，因此，這可能意味著前后兩次運行之間的GPU基本溫度可能不同，并可能稍微影響到運算時間。
結(jié)論
可以看出，選擇2的n次冪或8的倍數(shù)作為batch大小在實踐中不會產(chǎn)生明顯差異。
然而，由于在實際使用中已成為約定俗成，選擇2的n次冪作為batch大小，的確可以幫助運算更簡單并且易于管理。
此外，如果你有興趣發(fā)表學(xué)術(shù)研究論文，選擇2的n次冪將使你的論文看上去不那么主觀。
盡管如此，R教授仍然認為，batch的最佳大小在很大程度上取決于神經(jīng)網(wǎng)絡(luò)架構(gòu)和損失函數(shù)。
例如，在最近使用相同ResNet架構(gòu)的研究項目中，他發(fā)現(xiàn)batch的最佳大小可以在16到256之間，具體取決于損失函數(shù)。
因此，R教授建議始終把調(diào)整batch大小，作為超參數(shù)優(yōu)化的一部分。
但是，如果你由于內(nèi)存限制而無法使用512作為batch大小，那么則不必降到256，首先考慮500即可。
作者Sebastian Raschka
Sebastian Raschka，是一名機器學(xué)習(xí)和 AI 研究員。
他在UW-Madison（威斯康星大學(xué)麥迪遜分校）擔(dān)任統(tǒng)計學(xué)助理教授，專注于深度學(xué)習(xí)和機器學(xué)習(xí)研究，同時也是Lightning AI的首席 AI 教育家。
另外他還寫過一系列用Python和Scikit-learn做機器學(xué)習(xí)的教材。
基準測試代碼鏈接：
https://github.com/rasbt/b3-basic-batchsize-benchmark
參考鏈接：
https://sebastianraschka.com/blog/2022/batch-size-2.html
相關(guān)閱讀
