比用Pytorch框架快200倍！0.76秒后，筆記本上的CNN就搞定了MNIST | 開(kāi)源
博雯 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI
在MNIST上進(jìn)行訓(xùn)練，可以說(shuō)是計(jì)算機(jī)視覺(jué)里的“Hello World”任務(wù)了。
而如果使用PyTorch的標(biāo)準(zhǔn)代碼訓(xùn)練CNN，一般需要3分鐘左右。
但現(xiàn)在，在一臺(tái)筆記本電腦上就能將時(shí)間縮短200多倍。
速度直達(dá)0.76秒！
那么，到底是如何僅在一次epoch的訓(xùn)練中就達(dá)到99%的準(zhǔn)確率的呢？
八步提速200倍
這是一臺(tái)裝有GeForce GTX 1660 Ti GPU的筆記本。
我們需要的還有Python3.x和Pytorch 1.8。
先下載數(shù)據(jù)集進(jìn)行訓(xùn)練，每次運(yùn)行訓(xùn)練14個(gè)epoch。
這時(shí)兩次運(yùn)行的平均準(zhǔn)確率在測(cè)試集上為99.185%，平均運(yùn)行時(shí)間為2min 52s ± 38.1ms。
接下來(lái)，就是一步一步來(lái)減少訓(xùn)練時(shí)間：
一、提前停止訓(xùn)練
在經(jīng)歷3到5個(gè)epoch，測(cè)試準(zhǔn)確率達(dá)到99%時(shí)就提前停止訓(xùn)練。
這時(shí)的訓(xùn)練時(shí)間就減少了1/3左右，達(dá)到了57.4s±6.85s。
二、縮小網(wǎng)絡(luò)規(guī)模，采用正則化的技巧來(lái)加快收斂速度
具體的，在第一個(gè)conv層之后添加一個(gè)2×2的最大采樣層（max pool layer），將全連接層的參數(shù)減少4倍以上。
然后再將2個(gè)dropout層刪掉一個(gè)。
這樣，需要收斂的epoch數(shù)就降到了3個(gè)以下，訓(xùn)練時(shí)間也減少到30.3s±5.28s。
三、優(yōu)化數(shù)據(jù)加載
使用data_loader.save_data()，將整個(gè)數(shù)據(jù)集以之前的處理方式保存到磁盤的一個(gè)pytorch數(shù)組中。
也就是不再一次一次地從磁盤上讀取數(shù)據(jù)，而是將整個(gè)數(shù)據(jù)集一次性加載并保存到GPU內(nèi)存中。
這時(shí)，我們只需要一次epoch，就能將平均訓(xùn)練時(shí)間下降到7.31s ± 1.36s。
四、增加Batch Size
將Batch Size從64增加到128，平均訓(xùn)練時(shí)間減少到4.66s ± 583ms。
五、提高學(xué)習(xí)率
使用Superconvergence來(lái)代替指數(shù)衰減。
在訓(xùn)練開(kāi)始時(shí)學(xué)習(xí)率為0，到中期線性地最高值（4.0），再慢慢地降到0。
這使得我們的訓(xùn)練時(shí)間下降到3.14s±4.72ms。
六、再次增加Batch Size、縮小縮小網(wǎng)絡(luò)規(guī)模
重復(fù)第二步，將Batch Size增加到256。
重復(fù)第四步，去掉剩余的dropout層，并通過(guò)減少卷積層的寬度來(lái)進(jìn)行補(bǔ)償。
最終將平均時(shí)間降到1.74s±18.3ms。
七、最后的微調(diào)
首先，將最大采樣層移到線性整流函數(shù)（ReLU）激活之前。
然后，將卷積核大小從3增加到5.
最后進(jìn)行超參數(shù)調(diào)整：
使學(xué)習(xí)率為0.01（默認(rèn)為0.001），beta1為0.7（默認(rèn)為0.9），bata2為0.9（默認(rèn)為0.999）。
到這時(shí)，我們的訓(xùn)練已經(jīng)減少到一個(gè)epoch，在762ms±24.9ms的時(shí)間內(nèi)達(dá)到了99.04%的準(zhǔn)確率。
“這只是一個(gè)Hello World案例”
對(duì)于這最后的結(jié)果，有人覺(jué)得司空見(jiàn)慣：
優(yōu)化數(shù)據(jù)加載時(shí)間，縮小模型尺寸，使用ADAM而不是SGD等等，都是常識(shí)性的事情。
我想沒(méi)有人會(huì)真的費(fèi)心去加速運(yùn)行MNIST，因?yàn)檫@是機(jī)器學(xué)習(xí)中的“Hello World”，重點(diǎn)只是像你展示最小的關(guān)鍵值，讓你熟悉這個(gè)框架——事實(shí)上3分鐘也并不長(zhǎng)吧。
而也有網(wǎng)友覺(jué)得，大多數(shù)人的工作都不在像是MNIST這樣的超級(jí)集群上。因此他表示：
我所希望的是工作更多地集中在真正最小化訓(xùn)練時(shí)間方面。
GitHub：
https://github.com/tuomaso/train_mnist_fast
參考鏈接：
[1]https://www.reddit.com/r/MachineLearning/comments/p1168k/p_training_cnn_to_99_on_mnist_in_less_than_1/
相關(guān)閱讀
華為北大等聯(lián)手打造的Transformer竟在CV領(lǐng)域超過(guò)了CNN：多項(xiàng)底層視覺(jué)任務(wù)達(dá)到SOTA
多種任務(wù)共享一個(gè)Transformer 模塊。
CNN更新?lián)Q代！性能提升算力減半，還即插即用
OctConv即插即用，無(wú)需修改原來(lái)的網(wǎng)絡(luò)架構(gòu)，也不用調(diào)整超參數(shù)，方便到家。
MNIST數(shù)據(jù)集重生！測(cè)試圖片增加到6萬(wàn)張，LeCun點(diǎn)贊
與標(biāo)準(zhǔn)的Pytorch MNIST數(shù)據(jù)加載器兼容。
