??最近在做某個(gè)視覺任務(wù)的模型訓(xùn)練，由于數(shù)據(jù)量比較少為了效果好一點(diǎn)，決定現(xiàn)在imagenet上pretrain一下骨干網(wǎng)絡(luò)。但是在訓(xùn)練的時(shí)候遇到了一個(gè)問題：每隔num_workers個(gè)iteration數(shù)據(jù)加載都很慢，通過查找資料和代碼搞清了這個(gè)問題。

背景

??設(shè)計(jì)了一個(gè)網(wǎng)絡(luò)做目標(biāo)檢測，骨干網(wǎng)絡(luò)是自己diy的因此沒有pretrain的模型。而目標(biāo)檢測的數(shù)據(jù)集比較小，為了把模型訓(xùn)的好一點(diǎn)決定現(xiàn)把骨干網(wǎng)絡(luò)搭一個(gè)分類頭做個(gè)分類模型，在ImageNet上面pretrain一下。于是乎下載了imagenet的數(shù)據(jù)集，訓(xùn)練集總共120多萬張圖，驗(yàn)證集5萬張，全部都按圖片兒存儲。自己寫了一個(gè)dataset，然后打包成一個(gè)dataloader做訓(xùn)練的循環(huán)。為了加快數(shù)據(jù)讀取速度，dataloader里面指定了num_workers=4。由于gpu顯存比較充裕，設(shè)置了很大的batch_size=384.

問題

??在訓(xùn)練的時(shí)候發(fā)現(xiàn)，每當(dāng)?shù)螖?shù)能整除4的時(shí)候數(shù)據(jù)加載的速度都很慢，否則就比較快。改變num_workers為不同的值都會出現(xiàn)類似的情況，每當(dāng)?shù)螖?shù)能被num_workers整除時(shí)，數(shù)據(jù)加載都很慢，否則就比較快。只有當(dāng)設(shè)置num_workers=0時(shí)即只通過主進(jìn)程加載數(shù)據(jù)時(shí)，每一迭代的耗時(shí)才比較均勻。比如，當(dāng)把num_workers設(shè)置為4時(shí)，抓下來的log:

train batch 1 load time 21.43259024620056s
train batch 2 load time 0.031423091888427734s
train batch 3 load time 0.004406452178955078
train batch 4 load time 0.004347562789916992
train batch 5 load time 18.13344931602478
train batch 6 load time 0.004399538040161133
train batch 7 load time 0.03353142738342285
train batch 8 load time 0.004467010498046875
train batch 9 load time 16.202253103256226
train batch 10 load time 0.8377358913421631
…

原因分析

??這個(gè)問題的原因其實(shí)比較簡單：數(shù)據(jù)加載的速度太慢了，dataloader加載一個(gè)batch的時(shí)間遠(yuǎn)遠(yuǎn)大于gpu消費(fèi)一個(gè)batch的時(shí)間，導(dǎo)致gpu大多時(shí)候都在等待。畫個(gè)圖來說明一下。
??首先構(gòu)建了一個(gè)dataset，里面裝著120多萬張ImageNet的訓(xùn)練集圖片，然后通過一個(gè)num_workers=4的dataloader從數(shù)據(jù)池里不停的打包batch。num_workers=4相當(dāng)于開啟了4個(gè)子進(jìn)程分別獨(dú)立的進(jìn)行batch的打包，此時(shí)GPU在焦急的等待，他說你們4個(gè)進(jìn)程小兄弟，任何一個(gè)打包完一個(gè)batch都立刻給我我好拿去前向傳播。由于資源調(diào)度等細(xì)微差別，4個(gè)打包batch的進(jìn)程速度基本相同但稍有區(qū)別。不妨假設(shè)某次打包，四個(gè)進(jìn)程同時(shí)開始，進(jìn)程1完成當(dāng)前batch打包需要100ms，進(jìn)程2完成當(dāng)前batch打包需要98ms，進(jìn)程3完成當(dāng)前batch打包需要101ms，進(jìn)程4完成當(dāng)前batch打包需要102ms。

??那么從循環(huán)結(jié)果的角度來看，相當(dāng)于需要等待98ms才能進(jìn)行第一次循環(huán)，進(jìn)程2辛辛苦苦打包好的一個(gè)batch被拿去消費(fèi)了，然后他就默默的再去進(jìn)行下一次打包了。而僅僅過了2ms進(jìn)程1就又打包好了一個(gè)batch供給一次gpu循環(huán)，再過了僅僅1ms進(jìn)行第三次循環(huán)，再過1ms進(jìn)行第4次循環(huán)。此時(shí)共進(jìn)行了4次循環(huán)，分別把4個(gè)進(jìn)程辛苦了大約100ms才打包好的數(shù)據(jù)全用掉了。雖然進(jìn)程2在第98ms就已經(jīng)開始進(jìn)行下一個(gè)batch的打包了，但是仍然需要再過大約100ms才能準(zhǔn)備好下一次。而且由于資源調(diào)度等方面的隨機(jī)原因，最早開始下一輪打包數(shù)據(jù)的進(jìn)程，不見得在這一輪仍然是第一個(gè)完成打包的。
??單個(gè)進(jìn)程打包一個(gè)batch需要大概100ms，4個(gè)進(jìn)程同時(shí)工作結(jié)合上進(jìn)程的調(diào)度相當(dāng)于dataloader打包好一個(gè)batch的等效時(shí)間為:100/4=25ms，可是gpu消費(fèi)數(shù)據(jù)的速度很快，比如只需要10ms。那么就會造成gpu在每個(gè)num_workers次循環(huán)都需要等待很長時(shí)間，看看時(shí)序圖就清楚了。

??此時(shí)從循環(huán)時(shí)序的角度來看：每當(dāng)?shù)螖?shù)被子進(jìn)程數(shù)整除時(shí)，迭代很慢因?yàn)楸粩?shù)據(jù)加載拖累；除此之外迭代都很快，因?yàn)間pu剛消費(fèi)完一個(gè)batch另外一個(gè)batch也已經(jīng)準(zhǔn)備好了，gpu直接拿去再次消費(fèi)沒有等待時(shí)間。

解決辦法

??根據(jù)上面討論的原因，不難得出以下幾種可能的解決辦法，最根本的邏輯就在于：要讓多個(gè)子進(jìn)程并行情況下的等效batch打包時(shí)間，恰好等于gpu的消費(fèi)batch的時(shí)間，此時(shí)既不會出現(xiàn)dataloader已經(jīng)打包好了batch等待gpu使用，也不會出現(xiàn)gpu算力已經(jīng)空閑等待dataloader打包batch ，可以充分的將整個(gè)pipeline的資源全部利用起來。

增加num_workers數(shù)值

??dataloader的打包batch等效時(shí)間理論上等于：單個(gè)進(jìn)程打包時(shí)間/進(jìn)程數(shù)量。因此增加num_workers數(shù)值，增加進(jìn)程數(shù)可以縮小數(shù)據(jù)打包和gpu消費(fèi)數(shù)據(jù)之間的時(shí)間差。但是需要注意的是進(jìn)程之間的調(diào)度也是有時(shí)間消耗的，因此并不是進(jìn)程數(shù)越多打包速度越快。當(dāng)打包時(shí)間和gpu消耗batch時(shí)間相差不大時(shí)，可以考慮通過這種方法優(yōu)化。

減小數(shù)據(jù)讀取時(shí)間

??進(jìn)程打包時(shí)間很長的一個(gè)重要原因是，我在做ImageNet預(yù)訓(xùn)練時(shí)原始數(shù)據(jù)是一張張的圖片兒，batch size設(shè)置為384是一個(gè)比較大的數(shù)，單個(gè)進(jìn)程打包一個(gè)batch需要進(jìn)行很多張圖片文件的打開、讀取和關(guān)閉操作，IO時(shí)間遠(yuǎn)遠(yuǎn)大于數(shù)據(jù)讀取和處理時(shí)間，因此可以考慮減小數(shù)據(jù)的IO操作次數(shù)提高數(shù)據(jù)讀取速度，例如類似于tensorflow把數(shù)據(jù)轉(zhuǎn)成tfrecord這種大尺寸的二進(jìn)制文件等。
??有關(guān)這個(gè)部分，可以看一下我寫的另外一篇講dali的文章，完美的解決了數(shù)據(jù)讀取的時(shí)間問題，gpu利用率可以達(dá)到幾乎100%。

??除了上面的2種，應(yīng)該還有其他辦法解決這個(gè)問題，這里就不多羅列了。水平有限，歡迎討論。

總結(jié)

以上是生活随笔為你收集整理的Pytorch采坑记录：每隔num_workers个iteration数据加载速度很慢的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。