高性能 PyTorch 训练 (3)：并行与异步优化

如前文所述，在 PyTorch 训练过程中，一个主要的性能瓶颈就是数据集的导入和预处理。

在一般情况下，PyTorch 的 DataLoader 是实现了 __iter__ 和 __next__ 方法的，在 for 循环中，临时来从硬盘中导入数据，进行预处理并送入显存。这是 PyTorch 的一般实践，可以保证训练稳定地进行而不会出错。但是另一方面，稳定就代表不够高效。以 CIFAR10 数据集为例，在正常的训练中，每次导入训练数据大概 64 至 128 张图片左右，每张图片的尺寸为 [3, 32, 32]。再加上预处理的时间，每个 batch 的导入就需要花费近数秒钟。而整个训练过程动辄需要导入 200k+ 的 batch，花费的额外时间可想而知。

如果我们充分从代码层面上减少数据导入所用到的时间，就可以充分利用硬件的性能，来提高整个训练过程的效率。

数据导入时的并行优化

默认情况下，一个 batch 的数据导入也是串行的。显然，这是非常不合理的，因为数据与数据之间没有任何的依赖关系。现在的主流 CPU 至少 4 个物理核心起步，好一点的服务器甚至有几百个物理核之多。这时候就需要在 DataLoader 实例化的时候进行定义。

DataLoader(..., pin_memory=True, num_workers=xxx, ...)

其中，pin_memory=True 表示将数据导入到 GPU 的固定显存中，压缩数据加载和转移的时间。而 num_workers 就很明显了，顾名思义，就是指导入进程时的子进程数。一般来说，建议设置为 CPU 的总线程数。

这样，数据就会以并行的方式加载到显存中以供调用。

数据导入时的异步优化

我们可以设想以下两种过程：

• 所有的过程都是同步的。在一个周期内，需要先等待时间开销为 t1 的数据加载完成，然后进行时间开销为 t2 的训练过程，总时间开销为 t1+t2；
• 异步地进行数据加载。在进行训练的同时，加载下一次训练所需要的数据，并等待训练完成。这样的总时间开销为 t2。

显然，对于数据量越大的任务，t1 会带来更大的影响。

实现这种异步加载的方式也很简单。我们知道，CPU 和 GPU 之间是类似于 C/S 架构的关系，即由 CPU 发出指令并进行由一个线程进行异步等待，而 GPU 收到指令后分配一定的资源进行运算，并返回给阻塞等待的线程。可以充分利用这种关系，通过异步的形式来实现整个的加速过程。

class CudaDataLoader:
    """ 异步预先将数据从 CPU 加载到 GPU 中 """

    def __init__(self, loader, device, queue_size=2):
        self.device = device
        self.queue_size = queue_size
        self.loader = loader

        self.load_stream = torch.cuda.Stream(device=device)
        self.queue = Queue(maxsize=self.queue_size)

        self.idx = 0
        self.worker = Thread(target=self.load_loop)
        self.worker.setDaemon(True)
        self.worker.start()

    def load_loop(self):
        """ 不断的将 cuda 数据加载到队列里 """
        # The loop that will load into the queue in the background
        while True:
            for i, sample in enumerate(self.loader):
                self.queue.put(self.load_instance(sample))

    def load_instance(self, sample):
        """ 将 batch 数据从 CPU 加载到 GPU 中 """
        if torch.is_tensor(sample):
            with torch.cuda.stream(self.load_stream):
                return sample.to(self.device, non_blocking=True)
        elif sample is None or type(sample) in (list, str):
            return sample
        elif isinstance(sample, dict):
            return {k: self.load_instance(v) for k, v in sample.items()}
        else:
            return [self.load_instance(s) for s in sample]

    def __iter__(self):
        self.idx = 0
        return self

    def __next__(self):
        # 加载线程意外退出了
        if not self.worker.is_alive() and self.queue.empty():
            self.idx = 0
            self.queue.join()
            self.worker.join()
            raise StopIteration
        # 一个 epoch 加载完了
        elif self.idx >= len(self.loader):
            self.idx = 0
            raise StopIteration
        # 下一个 batch
        else:
            out = self.queue.get()
            self.queue.task_done()
            self.idx += 1
        return out

    def next(self):
        return self.__next__()

    def __len__(self):
        return len(self.loader)

    @property
    def sampler(self):
        return self.loader.sampler

    @property
    def dataset(self):
        return self.loader.dataset

这样，多 CPU 线程加载数据然后异步转移到 GPU 上整一条链路都进行了优化。

我们还可以继续进行一定的优化。在使用中，我发现在每个 epoch 开始加载数据的时候，DataLoader 还是会重新初始化 CPU 线程以及队列之类的，凭空消耗了相当多的垃圾时间。可以直接使用上次的就好。

class _RepeatSampler(object):
    """ 一直repeat的sampler """

    def __init__(self, sampler):
        self.sampler = sampler

    def __iter__(self):
        while True:
            yield from iter(self.sampler)

class MultiEpochsDataLoader(torch.utils.data.DataLoader):
    """ 多 epoch 训练时，DataLoader 对象不用重新建立线程和 batch_sampler 对象，以节约每个 epoch 的初始化时间 """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler))
        self.iterator = super().__iter__()

    def __len__(self):
        return len(self.batch_sampler.sampler)

    def __iter__(self):
        for i in range(len(self)):
            yield next(self.iterator)

综上，真正使用时，只需要上面的代码，重新封装 DataLoader 即可。

loader = MultiEpochsDataLoader(data_set, batch_size=128, 
                               shuffle=True, num_workers=32, 
                               pin_memory=True)
if torch.cuda.is_available():
    loader = CudaDataLoader(loader, 'cuda')

for image in loader:
    train_model(image)

在我的实际使用中，每个 epoch 可以加速接近 40%。