pytorch鎖死在dataloader(訓(xùn)練時卡死)

1.問題描述

2.解決方案

（1）Dataloader里面不用cv2.imread進(jìn)行讀取圖片，用cv2.imread還會帶來一系列的不方便，比如不能結(jié)合torchvision進(jìn)行數(shù)據(jù)增強(qiáng)，所以最好用PIL 里面的Image.open來讀圖片。（并不適用本例）

（2）將DataLoader 里面的參變量num_workers設(shè)置為0，但會導(dǎo)致數(shù)據(jù)的讀取很慢，拖慢整個模型的訓(xùn)練。（并不適用本例）

（3）如果用了cv2.imread，不想改代碼的，那就加兩條語句，來關(guān)閉Opencv的多線程：cv2.setNumThreads(0)和cv2.ocl.setUseOpenCL(False)。加了這兩條語句之后，并不影響模型的訓(xùn)練。（并不適用本例）

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

參考建議

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

如果GPU在t秒內(nèi)運(yùn)行每個迭代，而每個dataloader worker加載/處理單個批處理需要N*t秒，那么您應(yīng)該將num_workers設(shè)置為至少N，以避免GPU停滯。當(dāng)然，系統(tǒng)中至少要有N個cpu。

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

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

您還應(yīng)該在_get_item__中優(yōu)化您的代碼，以便每個worker在較短的時間內(nèi)完成其批處理。請確保worker完成批處理的時間不受從磁盤讀取訓(xùn)練數(shù)據(jù)的時間(特別是當(dāng)您從網(wǎng)絡(luò)存儲中讀取數(shù)據(jù)時)或網(wǎng)絡(luò)帶寬(當(dāng)您從網(wǎng)絡(luò)磁盤讀取數(shù)據(jù)時)的影響。如果您的數(shù)據(jù)集很小，并且您有足夠的RAM，那么可以考慮將數(shù)據(jù)集移動到RAM(或/tmpfs)中，并從那里讀取數(shù)據(jù)以進(jìn)行快速訪問。對于Kubernetes，您可以創(chuàng)建一個RAM磁盤(在Kubernetes中搜索emptyDir)。

如果你已經(jīng)優(yōu)化了你的_get_item__代碼，并確保磁盤訪問/網(wǎng)絡(luò)訪問不是罪魁禍?zhǔn)祝匀粫霈F(xiàn)問題，你將需要請求更多的cpu(為了一個Kubernetes pod)，或者將你的GPU移動到擁有更多cpu的機(jī)器上。

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

你也可以考慮離線做一些預(yù)處理，減輕每個worker的負(fù)擔(dān)。例如，如果每個worker正在讀取一個wav文件并計算音頻文件的譜圖，那么可以考慮離線預(yù)先計算譜圖，只從工作者的磁盤中讀取計算的譜圖。這將減少每個worker的工作量。

你也可以考慮將dataloader里的設(shè)置pin_memory=False。

補(bǔ)充：pytorch加載訓(xùn)練數(shù)據(jù)集dataloader操作耗費(fèi)時間太久，該如何解決？

筆者在使用pytorch加載訓(xùn)練數(shù)據(jù)進(jìn)行模型訓(xùn)練的時候，發(fā)現(xiàn)數(shù)據(jù)加載需要耗費(fèi)太多時間，該如何縮短數(shù)據(jù)加載的時間消耗呢？經(jīng)過查詢相關(guān)文檔，

總結(jié)實際操作過程如下：

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

2、dataloader函數(shù)中增加num_workers參數(shù)，該參數(shù)表示加載數(shù)據(jù)的線程數(shù)，建議設(shè)置為該系統(tǒng)中的CPU核心數(shù)，若CPU很強(qiáng)勁，而且內(nèi)存很大，也可以考慮將該數(shù)值設(shè)置的更大一些。

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達(dá)到了iter（Dataset）的讀取并行，但是沒有實現(xiàn)在GPU運(yùn)算時異步讀取數(shù)據(jù)，可以考慮使用non_blocking實現(xiàn)。

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才會有效。

以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持腳本之家。