pytorch鎖死在dataloader(訓練時卡死)

Posted on by

1.問題描述

2.解決方案

（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)。加瞭這兩條語句之後，並不影響模型的訓練。（並不適用本例）

（4）這種情況應該是屬於pytorch多線程鎖死，在github上看到有該問題，但是沒有解決的。

參考建議

首先確保num_works數量低於CPU數量(如果使用Kubernetes，則設置為pod)，但是設置得足夠高，使數據隨時可以用於下一次迭代。

如果GPU在t秒內運行每個迭代，而每個dataloader worker加載/處理單個批處理需要N*t秒，那麼您應該將num_workers設置為至少N，以避免GPU停滯。當然，系統中至少要有N個cpu。

不幸的是，如果Dataloader使用任何使用K個線程的庫，那麼生成的進程數量就會變成num_workersK = NK。這可能比計算機中的cpu數量大得多。這會使pod節流，而Dataloader會變得非常慢。這可能導致Dataloader不返回批處理每t秒，導致GPU暫停。

避免K個線程的一種方法是通過OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python train.py調用主腳本。這就限制瞭每個Dataloader工作程序隻能使用一個線程，從而避免瞭使機器不堪重負。你仍然需要有足夠的num_workers來滿足GPU的需要。

您還應該在_get_item__中優化您的代碼，以便每個worker在較短的時間內完成其批處理。請確保worker完成批處理的時間不受從磁盤讀取訓練數據的時間(特別是當您從網絡存儲中讀取數據時)或網絡帶寬(當您從網絡磁盤讀取數據時)的影響。如果您的數據集很小，並且您有足夠的RAM，那麼可以考慮將數據集移動到RAM(或/tmpfs)中，並從那裡讀取數據以進行快速訪問。對於Kubernetes，您可以創建一個RAM磁盤(在Kubernetes中搜索emptyDir)。

如果你已經優化瞭你的_get_item__代碼，並確保磁盤訪問/網絡訪問不是罪魁禍首，但仍然會出現問題，你將需要請求更多的cpu(為瞭一個Kubernetes pod)，或者將你的GPU移動到擁有更多cpu的機器上。

另一個選項是減少batch_size，這樣每個worker要做的工作就會減少，並且可以更快地完成預處理。後一種選擇在某些情況下是不可取的，因為會有空閑的GPU內存不被利用。

你也可以考慮離線做一些預處理，減輕每個worker的負擔。例如，如果每個worker正在讀取一個wav文件並計算音頻文件的譜圖，那麼可以考慮離線預先計算譜圖，隻從工作者的磁盤中讀取計算的譜圖。這將減少每個worker的工作量。

你也可以考慮將dataloader裡的設置pin_memory=False。

補充：pytorch加載訓練數據集dataloader操作耗費時間太久，該如何解決？

筆者在使用pytorch加載訓練數據進行模型訓練的時候，發現數據加載需要耗費太多時間，該如何縮短數據加載的時間消耗呢？經過查詢相關文檔，

總結實際操作過程如下：

1、盡量將jpg等格式的文件保存為bmp文件，可以降低解碼時間；

2、dataloader函數中增加num_workers參數，該參數表示加載數據的線程數，建議設置為該系統中的CPU核心數，若CPU很強勁，而且內存很大，也可以考慮將該數值設置的更大一些。

train_loader=torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True)

修改為：

train_loader=torch.utils.data.DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle=True,num_workers=multiprocessing.cpu_count())

雖然使用dataloader達到瞭iter（Dataset）的讀取並行，但是沒有實現在GPU運算時異步讀取數據，可以考慮使用non_blocking實現。

dataloader = data.Dataloader(dataset, batch_size = batch_size, num_workers = workers)
for epoch in range(epochs):
    for batch_idx, (images, labels) in enumerate(dataloader):
        images = images.to(device)
        labels = labels.to(device)

改為：

dataloader = data.Dataloader(dataset, batch_size = batch_size, num_workers = workers, pin_memory = True)
for epoch in range(epochs):
    for batch_idx, (images, labels) in enumerate(dataloader):
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

需要註意的是：隻有pin_memory=True並且num_workers>0時non_blocking才會有效。

以上為個人經驗，希望能給大傢一個參考，也希望大傢多多支持WalkonNet。

推薦閱讀：