【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 等优化过的库,多多利用向量化代码,提升代码运行效率;
  • 尽量缩减数据大小,不要传输无用信息。

你可能感兴趣的:(pytorch,pytorch)