PyTorch 提高速度

原始文档:https://www.yuque.com/lart/ugkv9f/ugysgn 声明: 大部分内容来自知乎和其他博客的分享, 这里只作为一个收集罗列. 欢迎给出更多建议.

知乎回答(欢迎点赞哦):

  • pytorch dataloader 数据加载占用了大部分时间, 各位大佬都是怎么解决的? - 人民艺术家的回答 - 知乎

  • 使用 pytorch 时, 训练集数据太多达到上千万张, Dataloader 加载很慢怎么办? - 人民艺术家的回答 - 知乎

    更新日志

  • 2019年11月29日：更新一些模型设计技巧和推理加速的内容，补充了下apex的一个介绍链接， 另外删了tfrecord，pytorch能用么？这个我记得是不能，所以删掉了（表示删掉:<）。

  • 2019年11月30日：补充MAC的含义，补充ShuffleNetV2的论文链接
  • 2019年12月02日：之前说的pytorch不能用tfrecord，今天看到https://www.zhihu.com/question/358632497下的一个回答，涨姿势了。
  • 2019年12月23日：补充几篇关于模型压缩量化的科普性文章。

  • 2022年7月3日：重排格式，调整了重参数化的部分。

    预处理提速

  • 尽量减少每次读取数据时的预处理操作, 可以考虑把一些固定的操作, 例如 resize , 事先处理好保存下来, 训练的时候直接拿来用。

  • 将预处理搬到 GPU 上加速。
    • Linux 可以使用[NVIDIA/DALI](https://github.com/NVIDIA/DALI)。
    • 使用基于 Tensor 的图像处理操作。

  • 选择优化得更好的数据增强库，例如Albumentations。

    IO 提速

  • mmcv 对数据的读取提供了比较高效且全面的支持：OpenMMLab：MMCV 核心组件分析(三): FileClient

    加快图片读取

  • opencv 一般要比 PIL 要快 。

    • 请注意，PIL的惰性加载的策略使得其看上去open要比opencv的imread要快，但是实际上那并没有完全加载数据。可以对open返回的对象调用其load()方法，从而手动加载数据，这时的速度才是合理的。
  • 对于 jpeg 读取, 可以尝试 jpeg4py。
  • 存 bmp 图(降低解码时间)。

  • 关于不同图像处理库速度的讨论：Python 的各种 imread 函数在实现方式和读取速度上有何区别？ - 知乎

    降低读取次数

    对于大规模的小文件读取, 建议转成单独的文件, 可以选择考虑 [TFRecord（Tensorflow）](https://github.com/Lyken17/Efficient-PyTorch#data-loader) , [recordIO（recordIO）](https://github.com/Lyken17/Efficient-PyTorch#data-loader) , [hdf5](https://github.com/Lyken17/Efficient-PyTorch#data-loader) , [pth](https://github.com/Lyken17/Efficient-PyTorch#data-loader) , [n5](https://github.com/Lyken17/Efficient-PyTorch#data-loader) , [lmdb](https://github.com/Lyken17/Efficient-PyTorch#data-loader) 等。

  • TFRecord ：https://github.com/vahidk/tfrecord

  • lmdb 数据库：

  • 如何给你 PyTorch 里的 Dataloader 打鸡血 - MKFMIKU 的文章 - 知乎

  • 给 pytorch 读取数据加速 - 体 hi 的文章 - 知乎

    借助内存

  • 直接载到内存里面。

    • 将图片读取后存到一个固定的容器对象中。
      • YoloV5 中的[--cache](https://github.com/ultralytics/yolov5/blob/19f33cbae29ac2127dd877b52e228c178dda6086/utils/dataloaders.py#L521-L534)。

  • 把内存映射成磁盘。

    借助固态

    机械硬盘换成 NVME 固态。参考自如何给你 PyTorch 里的 Dataloader 打鸡血 - MKFMIKU 的文章 - 知乎

    训练策略

    更低的精度

    在训练中使用低精度( FP16 甚至 INT8 、二值网络、三值网络)表示取代原有精度( FP32 )表示。
    可以节约一定的显存并提速, 但是要小心数值溢出。

  • 混合精度训练的介绍文章：

  • [NVIDIA/Apex](https://github.com/nvidia/apex)提供的混合精度支持。

  • PyTorch1.6 开始提供的[torch.cuda.amp](https://pytorch.org/docs/stable/notes/amp_examples.html)以支持混合精度。

    更大的批次

    更大的批次在固定训练周期的情况下往往会带来更短的训练时间。但是大的批次面临着超参数的设置、显存占用问题等诸多考量，这又是另一个备受关注的领域了。

  • 超参数设置

    • Accurate, large minibatch SGD: training imagenet in 1 hour，论文

  • 优化显存占用

  • 在训练循环之前设置torch.backends.cudnn.benchmark = True可以加速计算。由于计算不同内核大小卷积的 cuDNN 算法的性能不同，自动调优器可以运行一个基准来找到最佳算法。当你的输入大小不经常改变时，建议开启这个设置。如果输入大小经常改变，那么自动调优器就需要太频繁地进行基准测试，这可能会损害性能。它可以将向前和向后传播速度提高 1.27x 到 1.70x。

  • 使用页面锁定内存，即在 DataLoader 中设定[pin_memory=True](https://pytorch.org/docs/stable/data.html#memory-pinning)。
  • 合适的 num_worker，细节讨论可见Pytorch 提速指南 - 云梦的文章 - 知乎。
  • optimizer.zero_grad(set_to_none=False)这里可以通过设置set_to_none=True来降低的内存占用，并且可以适度提高性能。但是这也会改变某些行为，具体可见文档。通过model.zero_grad()或optimizer.zero_grad()将对所有参数执行memset，并通过读写操作更新梯度。但是，将梯度设置为None将不会执行memset，并且将使用“只写”操作更新梯度。因此，设置梯度为None更快。
  • 反向传播期间设定使用eval模式并使用torch.no_grad关闭梯度计算。
  • 可以考虑使用channels_last的内存格式。

  • 用[DistributedDataParallel](https://pytorch.org/docs/stable/notes/cuda.html#use-nn-parallel-distributeddataparallel-instead-of-multiprocessing-or-nn-dataparallel)代替[DataParallel](https://pytorch.org/docs/stable/notes/cuda.html#use-nn-parallel-distributeddataparallel-instead-of-multiprocessing-or-nn-dataparallel)。对于多 GPU 来说，即使只有单个节点，也总是优先使用 DistributedDataParallel而不是 DataParallel ，因为 DistributedDataParallel 应用于多进程，并为每个 GPU 创建一个进程，从而绕过 Python 全局解释器锁(GIL)并提高速度。

    模型搭建

  • 不要初始化任何用不到的变量，因为 PyTorch 的初始化和 forward 是分开的，他不会因为你不去使用，而不去初始化。

  • [@torch.jit.script](https://pytorch.org/docs/stable/generated/torch.jit.script.html#torch.jit.script)，使用 PyTroch JIT 将逐点运算融合到单个 CUDA kernel 上。
  • 在使用混合精度的 FP16 时，对于所有不同架构设计，设置尺寸为 8 的倍数。

  • BN 之前的卷积层可以去掉 bias。因为在数学上，bias 可以通过 BN 的均值减法来抵消。我们可以节省模型参数、运行时的内存。

    数据操作

  • 将 batch size 设置为 8 的倍数，最大化 GPU 内存的使用。

  • GPU 上尽可能执行 NumPy 风格的操作。
  • 使用del释放内存占用。
  • 避免不同设备之间不必要的数据传输。
  • 创建张量的时候，直接指定设备，而不要创建后再传输到目标设备上。
  • 使用[torch.from_numpy(ndarray)](https://pytorch.org/docs/stable/generated/torch.from_numpy.html#torch-from-numpy)或者[torch.as_tensor(data, dtype=None, device=None)](https://pytorch.org/docs/stable/generated/torch.as_tensor.html#torch-as-tensor)，这可以通过共享内存而避免重新申请空间，具体使用细节和注意事项可参考对应文档。如果源设备和目标设备都是 CPU，torch.from_numpy和torch.as_tensor不会拷贝数据。如果源数据是 NumPy 数组，使用torch.from_numpy更快。如果源数据是一个具有相同数据类型和设备类型的张量，那么torch.as_tensor可以避免拷贝数据，这里的数据可以是 Python 的 list， tuple，或者张量。

  • 使用非阻塞传输，即设定non_blocking=True。这会在可能的情况下尝试异步转换，例如，将页面锁定内存中的 CPU 张量转换为 CUDA 张量。

    优化器设置

  • 将模型参数存放到一块连续的内存中，从而减少optimizer.step()的时间。

    • [contiguous_pytorch_params](https://github.com/PhilJd/contiguous_pytorch_params)

  • 使用 APEX 中的fused building blocks

    模型设计

    CNN

  • ShuffleNetV2，论文。

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

      Vision Transformer

  • TRT-ViT: TensorRT-oriented Vision Transformer，论文，解读。

    • stage-level：Transformer block 适合放置到模型的后期，这可以最大化效率和性能的权衡。
    • stage-level：先浅后深的 stage 设计模式可以提升性能。
    • block-level：Transformer 和 BottleNeck 的混合 block 要比单独的 Transformer 更有效。
    • block-level：先全局再局部的 block 设计模式有助于弥补性能问题。

      通用思路

  • 降低复杂度: 例如模型裁剪和剪枝, 减少模型层数和参数规模。

  • 改模型结构: 例如模型蒸馏, 通过知识蒸馏方法来获取小模型。

    推理加速

    更低的精度

    在推理中使用低精度( FP16 甚至 INT8 、二值网络、三值网络)表示取代原有精度( FP32 )表示。

  • TensorRT 是 NVIDIA 提出的神经网络推理(Inference)引擎, 支持训练后 8BIT 量化, 它使用基于交叉熵的模型量化算法, 通过最小化两个分布的差异程度来实现

  • Pytorch1.3 开始已经支持量化功能, 基于 QNNPACK 实现, 支持训练后量化, 动态量化和量化感知训练等技术
  • 另外 Distiller 是 Intel 基于 Pytorch 开源的模型优化工具, 自然也支持 Pytorch 中的量化技术
  • 微软的 NNI 集成了多种量化感知的训练算法, 并支持 PyTorch/TensorFlow/MXNet/Caffe2 等多个开源框架

更多细节可参考有三 AI:【杂谈】当前模型量化有哪些可用的开源工具?。

算子融合