pytorch dataloader num_workers引數設定導致訓練阻塞

ShellCollector發表於2020-10-02

參考:https://blog.csdn.net/qian1122221/article/details/88579684

問題描述:
最近在用RFBnet (原始碼是pytorch的)訓練RSNA的比賽資料,除了要修改一點程式碼支援RSNA的資料集外(打算後續再寫個部落格),發現在使用dataloader讀取資料時,如果設定num_workers為0,也就是用主程式讀取資料,模型訓練程式執行正常。如果設定num_workers為其他任何一個大於0的整數,也就是使用子程式讀取資料時,訓練程式會卡住,卡在訓練之前,GPU使用率和視訊記憶體佔用率都為0。

 

解決過程:
由於是多程式就,用print除錯大法,定位到是一行opencv程式碼的問題,在dataloader子類的__getitem__方法裡面呼叫了

image_array = cv2.cvtColor(image_array, cv2.COLOR_GRAY2BGR)
所有子程式的程式碼都卡在這裡了。之前也有遇到過類似的問題,python多程式呼叫opencv的程式碼會出現問題。於是就用numpy的concate替代了這個方法

image_array = np.concatenate([image_array,image_array,image_array],axis=2)
繞過這個問題就可以正常訓練了。

 

問題探索:
其實這個問題,我在另外一臺機器上是沒有遇到了,該機器是python3.6+opencv3.4.2

遇到問題的環境是docker環境,python3.5+opencv3.2

我感覺跟opencv的版本可能有關係,等後續訓練完成之後把docker裡面opencv的版本升級到3.4.2再試一下

相關文章