PyTorch 訓練時在第一個 epoch 的最後一個 batch 卡死

- 問題描述

使用 PyTorch 框架訓練模型，訓練第一個 epoch 時，在最後一個 batch 處卡死，卡了一天都沒有動彈，而 CPU 和 GPU 都處於正常運行的狀態，程序也沒有報錯，並且之前訓練一直都是正常的。最終，只能通過 Ctrl+C 強制性暫停。如下圖所示。

- 可能的原因

1. 搜索文章發現，有人出現這種問題是和 cv2.imread 有關，用 OpenCV 的接口進行數據讀取，而沒有用 PIL，導致出現 OpenCV與Pytorch互鎖的問題，關閉OpenCV的多線程即可解決問題^{1 2}。
2. 但是我們的代碼中並沒有使用 OpenCV，通過 Debug 發現，出現問題的時候，最後一個 batch =1，而我們使用的是四塊 GPU 進行訓練，原因就在此。

- 解決方法

Pytorch 的多 GPU 處理接口是 torch.nn.DataParallel(module, device_ids)，具體並行機制可參考：Pytorch多GPU並行處理。
該接口還要求輸入數據的 batch 數量要不小於所指定的 GPU 數量。另根據官網的解釋和註釋 (The batch size should be larger than the number of GPUs used.)，batch的數量會均分到每塊GPU上進行處理，因此要保證一個整數的關係。

故，一定要注意在使用多塊 GPU 訓練時，避免出現batch 數量小於所指定的 GPU 數量的情況。

多 GPU 並行訓練中其他比較坑的地方可參考：pytorch多GPU數據並行模式 踩坑指南，和nn.DataParallel說再見。

---

1. Pytorch dataloader在加載最後一個batch時卡死 ↩︎

2. 在PyTorch訓練一個epoch時，模型不能接着訓練，Dataloader卡死 ↩︎