Pytorch 如何加速Dataloader提升數(shù)據(jù)讀取速度
在利用DL解決圖像問題時，影響訓練效率最大的有時候是GPU，有時候也可能是CPU和你的磁盤。
很多設(shè)計不當?shù)娜蝿?wù)，在訓練神經(jīng)網(wǎng)絡(luò)的時候，大部分時間都是在從磁盤中讀取數(shù)據(jù)，而不是做 Backpropagation 。
這種癥狀的體現(xiàn)是使用 Nividia-smi 查看 GPU 使用率時，Memory-Usage 占用率很高，但是 GPU-Util 時常為 0% ，如下圖所示：
如何解決這種問題呢？
在 Nvidia 提出的分布式框架 Apex 里面，我們在源碼里面找到了一個簡單的解決方案：
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ù)據(jù)，
那么對我們自己的訓練代碼只需要做下面的改造：
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): # 訓練代碼 #-------------升級后--------- data, label = prefetcher.next() iteration = 0 while data is not None: iteration += 1 # 訓練代碼 data, label = prefetcher.next()
這樣子我們的 Dataloader 就像打了雞血一樣提高了效率很多，如下圖：
當然，最好的解決方案還是從硬件上，把讀取速度慢的機械硬盤換成 NVME 固態(tài)吧～
補充：Pytorch設(shè)置多線程進行dataloader時影響GPU運行
使用PyTorch設(shè)置多線程（threads）進行數(shù)據(jù)讀取時，其實是假的多線程，他是開了N個子進程（PID是連續(xù)的）進行模擬多線程工作。
以載入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是否運行該程序
查看運行在gpu上的所有程序：
fuser -v /dev/nvidia*
如果沒有返回，則該程序并沒有在GPU上運行
指定GPU運行
將num_workers改成0即可
以上為個人經(jīng)驗，希望能給大家一個參考，也希望大家多多支持本站。
版權(quán)聲明：本站文章來源標注為YINGSOO的內(nèi)容版權(quán)均為本站所有，歡迎引用、轉(zhuǎn)載，請保持原文完整并注明來源及原文鏈接。禁止復制或仿造本網(wǎng)站，禁止在非maisonbaluchon.cn所屬的服務(wù)器上建立鏡像，否則將依法追究法律責任。本站部分內(nèi)容來源于網(wǎng)友推薦、互聯(lián)網(wǎng)收集整理而來，僅供學習參考，不代表本站立場，如有內(nèi)容涉嫌侵權(quán)，請聯(lián)系alex-e#qq.com處理。
關(guān)注官方微信