优化Pytorch的数据加载

背景

在利用深度学习解决图像问题时，影响训练效率最大的有时候是GPU，有时候也可能是CPU和你的磁盘。 很多设计不当的任务，在训练神经网络的时候，大部分时间都是在从磁盘中读取数据，而不是做 Backpropagation 。这种症状的体现是使用 Nividia-smi 查看 GPU 使用率时，Memory-Usage 占用率很高，但是 GPU-Util 时常为 0% ，如下图所示：

nvidia-smi

解决方案

如何解决这种问题呢？在 Nvidia 提出的分布式框架 Apex 里面，我们在源码里面找到了一个简单的解决方案： NVIDIA/apex

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 就像打了鸡血一样提高了效率很多，如下图：

nvidia-smi2

其他方案

预处理提速

  • 尽量减少每次读取数据时的预处理操作，可以考虑把一些固定的操作，例如 resize ，事先处理好保存下来，训练的时候直接拿来用
  • Linux上将预处理搬到GPU上加速：NVIDIA/DALI

IO提速

  • 使用更快的图片处理库
  •    opencv一般要比 PIL 要快
  •    对于jpeg读取，可以尝试 jpeg4py
  •    存bmp图（降低解码时间）
  • 小图拼起来存放（降低读取次数，对于大规模的小文件读取，建议转成单独的文件，可以选择的格式：TFRecord（Tensorflow）、recordIO（recordIO）、hdf5、 pth、n5、lmdb 等等 Efficient-PyTorch
  •   TFRecord
  •   借助lmdb数据库：     Image2LMDB     LMDB     PySODToolBox
  • 借助内存：直接载到内存里面，或者把把内存映射成磁盘
  • 借助固态：把读取速度慢的机械硬盘换成 NVME 固态

训练策略

  • 在训练中使用低精度（FP16甚至INT8、二值网络）表示取代原有精度（FP32）表示
  • NVIDIA/Apex使用
  • torch.backends.cudnn.benchmark = True
  • Do numpy-like operations on the GPU wherever you can
  • Free up memory using del
  • Avoid unnecessary transfer of data from the GPU
  • Use pinned memory, and use non_blocking=False to parallelize data transfer and GPU number crunching

模型设计

来自ShuffleNetV2的结论： 内存访问消耗时间，memory access cost 缩写为 MAC。

  • 卷积层输入输出通道一致：卷积层的输入和输出特征通道数相等时MAC最小，此时模型速度最快
  • 减少卷积分组：过多的group操作会增大MAC，从而使模型速度变慢
  • 减少模型分支：模型中的分支数量越少，模型速度越快
  • 减少element-wise操作：element-wise操作所带来的时间消耗远比在FLOPs上的体现的数值要多，因此要尽可能减少element-wise操作（depthwise convolution也具有低FLOPs、高MAC的特点）

其他：

  • 降低复杂度：例如模型裁剪和剪枝，减少模型层数和参数规模
  • 改模型结构：例如模型蒸馏，通过知识蒸馏方法来获取小模型

本文来源：

results matching ""

    No results matching ""