快捷導(dǎo)航

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

更新時(shí)間：2021年05月28日 09:34:02 作者：Totoro-wen

這篇文章主要介紹了pytorch鎖死在dataloader(訓(xùn)練時(shí)卡死)，具有很好的參考價(jià)值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教

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ù)的讀取很慢，拖慢整個(gè)模型的訓(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ù)隨時(shí)可以用于下一次迭代。

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

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

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

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

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

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

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

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

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

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

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

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

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）的讀取并行，但是沒有實(shí)現(xiàn)在GPU運(yùn)算時(shí)異步讀取數(shù)據(jù)，可以考慮使用non_blocking實(shí)現(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時(shí)non_blocking才會有效。

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

您可能感興趣的文章: