【pytorch】给训练踩踩油门-- Pytorch 加速数据读取

本文转载于 https://zhuanlan.zhihu.com/p/80695364

训练模型的时候有时候会发现显卡的占用一直跑不满，会很浪费，往往是因为IO瓶颈导致的训练速度降低。
本文可以从以下几个方面进行对模型加速：

一, prefetch_generator

使用 prefetch_generator 库在后台加载下一 batch 的数据。

安装：

pip install prefetch_generator

使用：

# 新建DataLoaderX类
from torch.utils.data import DataLoader
from prefetch_generator import BackgroundGenerator

class DataLoaderX(DataLoader):

    def __iter__(self):
        return BackgroundGenerator(super().__iter__())

然后用 DataLoaderX 替换原本的 DataLoader

提速原因：

原本 PyTorch 默认的 DataLoader 会创建一些 worker 线程来预读取新的数据，但是除非这些线程的数据全部都被清空，这些线程才会读下一批数据。
使用 prefetch_generator，我们可以保证线程不会等待，每个线程都总有至少一个数据在加载。

二, data_prefetcher

使用 data_prefetcher 新开 cuda stream 来拷贝 tensor 到 gpu。

class DataPrefetcher():
    def __init__(self, loader, opt):
        self.loader = iter(loader)
        self.opt = opt
        self.stream = torch.cuda.Stream()
        # 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.batch = next(self.loader)
        except StopIteration:
            self.batch = None
            return
        with torch.cuda.stream(self.stream):
            for k in self.batch:
                if k != 'meta':
                    self.batch[k] = self.batch[k].to(device=self.opt.device, 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()

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        batch = self.batch
        self.preload()
        return batch

# ----改造前----
for iter_id, batch in enumerate(data_loader):
    if iter_id >= num_iters:
        break
    for k in batch:
        if k != 'meta':
            batch[k] = batch[k].to(device=opt.device, non_blocking=True)
    run_step()
    
# ----改造后----
prefetcher = DataPrefetcher(data_loader, opt)
batch = prefetcher.next()
iter_id = 0
while batch is not None:
    iter_id += 1
    if iter_id >= num_iters:
        break
    run_step()
    batch = prefetcher.next()

提速原因：
默认情况下，PyTorch 将所有涉及到 GPU 的操作（比如内核操作，cpu->gpu，gpu->cpu）都排入同一个 stream（default stream）中，并对同一个流的操作序列化，它们永远不会并行。要想并行，两个操作必须位于不同的 stream 中。
而前向传播位于 default stream 中，因此，要想将下一个 batch 数据的预读取（涉及 cpu->gpu）与当前 batch 的前向传播并行处理，就必须：
（1） cpu 上的数据 batch 必须 pinned;
（2）预读取操作必须在另一个 stream 上进行
上面的 data_prefetcher 类满足这两个要求。注意 dataloader 必须设置 pin_memory=True 来满足第一个条件

三, 把内存当硬盘

把数据放内存里，降低 io 延迟。

sudo mount tmpfs /path/to/your/data -t tmpfs -o size=30G

然后把数据放挂载的目录下，即可。

• size 指定的是 tmpfs 动态大小的上限，实际大小根据实际使用情况而定；
• 数据不一定放在物理内存中，系统根据情况，有可能放在 swap 的页面，swap 一般是在系统盘；
• 重启或者断电后数据全部清空。

如果想系统启动时自动挂载，可以编辑 /etc/fstab，在最后添加如下内容：

mount tmpfs in /tmp/
tmpfs /tmp tmpfs size=30G 0 0

四, 设置num_worker

DataLoader 的 num_worker 如果设置太小，则不能充分利用多线程提速，如果设置太大，会造成线程阻塞，或者撑爆内存，反而导致训练变慢甚至程序崩溃。

他的大小和具体的硬件和软件都有关系，所以没有一个统一的标准，可以通过一些简单的实验来确定。

我的经验是设置成 cpu 的核心数或者 gpu 的数量比较合适。

五、 优化数据预处理

主要有两个方面：

• 尽量简化预处理的操作，使用 numpy、opencv 等优化过的库，多多利用向量化代码，提升代码运行效率；
• 尽量缩减数据大小，不要传输无用信息。