在PyTorch訓練一個epoch時,模型不能接着訓練,Dataloader卡死

    筆者在訓練模型的時候,突然偶遇這個問題,即訓練一個epoch時,模型不能接着訓練,只能通過Ctrl+C強制性暫停,見下圖:

    Ctrl+C之後呈現的信息表明,這個bug是和多線程有關係

    經過筆者實驗,目前有三種可靠的解決方式

    1).Dataloader裏面不用cv2.imread進行讀取圖片,用cv2.imread還會帶來一系列的不方便,比如不能結合torchvision進行數據增強,所以最好用PIL 裏面的Image.open來讀圖片

    2).將DataLoader 裏面的參變量num_workers設置爲0,但會導致數據的讀取很慢,拖慢整個模型的訓練

    3).如果用了cv2.imread,也懶得改了,那就加兩條語句,來關閉Opencv的多線程:cv2.setNumThreads(0)和cv2.ocl.setUseOpenCL(False)。加了這兩條語句之後,並不影響模型的訓練速度,根據筆者觀察,速度相反還變快了,真神奇呀。

    綜上所述:如果遇到此問題,建議選擇方法1和方法3來解決這個問題,因爲不影響模型訓練速度。

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