在利用DL解決圖像問題時(shí),影響訓(xùn)練效率最大的有時(shí)候是GPU,有時(shí)候也可能是CPU和你的磁盤。
很多設(shè)計(jì)不當(dāng)?shù)娜蝿?wù),在訓(xùn)練神經(jīng)網(wǎng)絡(luò)的時(shí)候,大部分時(shí)間都是在從磁盤中讀取數(shù)據(jù),而不是做 Backpropagation 。
這種癥狀的體現(xiàn)是使用 Nividia-smi 查看 GPU 使用率時(shí),Memory-Usage 占用率很高,但是 GPU-Util 時(shí)常為 0% ,如下圖所示:

如何解決這種問題呢?
在 Nvidia 提出的分布式框架 Apex 里面,我們在源碼里面找到了一個(gè)簡單的解決方案:
https://github.com/NVIDIA/apex/blob/f5cd5ae937f168c763985f627bbf850648ea5f3f/examples/imagenet/main_amp.py#L256
class data_prefetcher():
def __init__(self, loader):
self.loader = iter(loader)
self.stream = torch.cuda.Stream()
self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
# With Amp, it isn't necessary to manually convert data to half.
# if args.fp16:
# self.mean = self.mean.half()
# self.std = self.std.half()
self.preload()
def preload(self):
try:
self.next_input, self.next_target = next(self.loader)
except StopIteration:
self.next_input = None
self.next_target = None
return
with torch.cuda.stream(self.stream):
self.next_input = self.next_input.cuda(non_blocking=True)
self.next_target = self.next_target.cuda(non_blocking=True)
# With Amp, it isn't necessary to manually convert data to half.
# if args.fp16:
# self.next_input = self.next_input.half()
# else:
self.next_input = self.next_input.float()
self.next_input = self.next_input.sub_(self.mean).div_(self.std)
我們能看到 Nvidia 是在讀取每次數(shù)據(jù)返回給網(wǎng)絡(luò)的時(shí)候,預(yù)讀取下一次迭代需要的數(shù)據(jù),
那么對我們自己的訓(xùn)練代碼只需要做下面的改造:
training_data_loader = DataLoader(
dataset=train_dataset,
num_workers=opts.threads,
batch_size=opts.batchSize,
pin_memory=True,
shuffle=True,
)
for iteration, batch in enumerate(training_data_loader, 1):
# 訓(xùn)練代碼
#-------------升級后---------
data, label = prefetcher.next()
iteration = 0
while data is not None:
iteration += 1
# 訓(xùn)練代碼
data, label = prefetcher.next()
這樣子我們的 Dataloader 就像打了雞血一樣提高了效率很多,如下圖:

當(dāng)然,最好的解決方案還是從硬件上,把讀取速度慢的機(jī)械硬盤換成 NVME 固態(tài)吧~
補(bǔ)充:Pytorch設(shè)置多線程進(jìn)行dataloader時(shí)影響GPU運(yùn)行
使用PyTorch設(shè)置多線程(threads)進(jìn)行數(shù)據(jù)讀取時(shí),其實(shí)是假的多線程,他是開了N個(gè)子進(jìn)程(PID是連續(xù)的)進(jìn)行模擬多線程工作。
以載入cocodataset為例
DataLoader
dataloader = torch.utils.data.DataLoader(COCODataset(config["train_path"],
(config["img_w"], config["img_h"]),
is_training=True),
batch_size=config["batch_size"],
shuffle=True, num_workers=32, pin_memory=True)
numworkers就是指定多少線程的參數(shù),原為32。
檢查GPU是否運(yùn)行該程序
查看運(yùn)行在gpu上的所有程序:
如果沒有返回,則該程序并沒有在GPU上運(yùn)行
指定GPU運(yùn)行
將num_workers改成0即可
以上為個(gè)人經(jīng)驗(yàn),希望能給大家一個(gè)參考,也希望大家多多支持腳本之家。
您可能感興趣的文章:- pytorch鎖死在dataloader(訓(xùn)練時(shí)卡死)
- pytorch Dataset,DataLoader產(chǎn)生自定義的訓(xùn)練數(shù)據(jù)案例
- 解決Pytorch dataloader時(shí)報(bào)錯(cuò)每個(gè)tensor維度不一樣的問題
- pytorch中DataLoader()過程中遇到的一些問題
- Pytorch dataloader在加載最后一個(gè)batch時(shí)卡死的解決
- pytorch DataLoader的num_workers參數(shù)與設(shè)置大小詳解
- pytorch 實(shí)現(xiàn)多個(gè)Dataloader同時(shí)訓(xùn)練