Pytorch 如何加速Dataloader提升數據讀取速度

Posted on by

在利用DL解決圖像問題時，影響訓練效率最大的有時候是GPU，有時候也可能是CPU和你的磁盤。

很多設計不當的任務，在訓練神經網絡的時候，大部分時間都是在從磁盤中讀取數據，而不是做 Backpropagation 。

這種癥狀的體現是使用 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 是在讀取每次數據返回給網絡的時候，預讀取下一次迭代需要的數據，

那麼對我們自己的訓練代碼隻需要做下面的改造：

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 固態吧～

補充：Pytorch設置多線程進行dataloader時影響GPU運行

使用PyTorch設置多線程（threads）進行數據讀取時，其實是假的多線程，他是開瞭N個子進程（PID是連續的）進行模擬多線程工作。

以載入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就是指定多少線程的參數，原為32。

檢查GPU是否運行該程序

查看運行在gpu上的所有程序：

fuser -v /dev/nvidia*

如果沒有返回，則該程序並沒有在GPU上運行

指定GPU運行

將num_workers改成0即可

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

推薦閱讀：