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

背景

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

問題

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

原因分析

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

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

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

解決辦法

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

增加num_workers數值

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

減小數據讀取時間

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

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

總結

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

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