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

背景

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

問題

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

原因分析

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

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

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

解決辦法

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

增加num_workers數(shù)值

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

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

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

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

總結

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

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。