pytorch dataloader num_workers參數設置導致訓練阻塞

問題描述:

最近在用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再試一下

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章