正文

Pytorch Dataloader加速

IP云
IP云V管理员/2 阅读
11/30
提示：扫一扫查出行
【扫一扫了解最新限行尾号】
复制提示

在进行多卡训练的时候，经常会出现GPU利用率上不来的情况，无法发挥硬件的最大实力。 造成这种现象最有可能的原因是，CPU生成数据的能力，已经跟不上GPU处理数据的能力。

方法一

常见的方法为修改Dataloader里面的线程数量，利用多线程技术提高数据生产能力，但是这种方法提速并不是特别明显。 

train_loader = DataLoader(dataset, batch_size,shuffle=True, num_worker=4)

而且windows机器上，num_worker大于0时，有时会出现卡死的情况，这应该是pytorch的bug，因此不是特别建议这种方法。

不过这种方法最简单，还是可以尝试一下更改线程数能否缓解你遇到的问题。nun_worker一般设置为处理器的物理线程数，不宜过大，因为会导致额外的线程开销。

方法二

本文主要介绍第二种方法，也就是Data Prefetcher，最早见于NVIDIA APEX。

这里我把代码抠出来了，删除掉了一些不必要的注释，可以将其复用到自己的项目里来。 

import torch

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)

        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)

            self.next_input = self.next_input.float()

            self.next_input = self.next_input.sub_(self.mean).div_(self.std)

    def next(self):

        torch.cuda.current_stream().wait_stream(self.stream)

        input = self.next_input

        target = self.next_target

        if input is not None:

            input.record_stream(torch.cuda.current_stream())

        if target is not None:

            target.record_stream(torch.cuda.current_stream())

        self.preload()

        return input, target

首先我们来看初始化函数，在初始化函数中，会直接调用preload，所以当这个对象初始化时，就会生成第一份的输入数据。

核心逻辑也就在预加载函数preload中，其中第13行是从原来的dataloader中取数，这一步和常规数据加载没有差别。有差别的是第19行，这里出现了Stream的概念。

一般来说，CUDA程序默认都运行在同一个Stream上，因此CPU->GPU，GPU->GPU以及GPU->CPU的一系列计算都是在同一个Stream里面串行运行的。 深度学习一般流程是先从dataloader中取数，这里是内存->CPU的运算，然后执行to_device操作，让数据从CPU->GPU，再是GPU->GPU的神经网络计算。

代码19行，使得data_prefetecher这个类是单独运行在一个Stream上的，因此它让数据加载和神经网络计算可以并行执行，也就加速了整体的运行速度。这样做带来的负面结果就是GPU同时在做两项任务，所以显存占用会增加。

这里不知道解释清楚没有，建议去看一下原作者的回答link

另外，重要的是，使用这个方法的时候一定要将Dataloader里面的pin_memory设置为True。

使用方法如下，非常简单，改造前是从dataloader里取数，改造后是将dataloader包在prefetecher里面，从prefetecher里面取数。 

train_loader = DataLoader(dataset, batch_size,shuffle=True, num_worker=4,pin_memory=True)

prefetcher = data_prefetcher(train_loader)

input, target = prefetcher.next()

while input is not None:

    ##

    前后向计算...

    ###

    input, target = prefetcher.next()

相关阅读

目录[+]