• 问答
  • 技术
  • 实践
  • 资源
使用 pytorch 时,训练集数据太多达到上千万张,Dataloader 加载很慢怎么办?
技术讨论
来源:人民艺术家 (知乎)


预处理提速

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

  • NVIDIA/DALI :https://github.com/NVIDIA/DALI


IO提速

使用更快的图片处理

  • opencv 一般要比 PIL 要快
  • 对于 jpeg 读取,可以尝试 jpeg4py
  • bmp 图(降低解码时间)


小图拼起来存放(降低读取次数)

对于大规模的小文件读取,建议转成单独的文件,可以选择的格式可以考虑:TFRecord(Tensorflow)recordIO(recordIO)hdf5pthn5lmdb 等等(https://github.com/Lyken17/Efficient-PyTorch#data-loader


预读取数据

  • 预读取下一次迭代需要的数据

【参考】

借助内存

  • 直接载到内存里面,或者把把内存映射成磁盘好了

【参考】

借助固态

  • 把读取速度慢的机械硬盘换成 NVME 固态吧~

【参考】


训练策略

低精度训练

代码层面

  • 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=True to parallelize data transfer and GPU number crunching

  • 文档:https://pytorch.org/docs/stable/nn.html#torch.nn.Module.to
  • 关于 non_blocking=True 的设定的一些介绍:Pytorch有什么节省显存的小技巧? - 陈瀚可的回答 - 知乎 https://www.zhihu.com/question/274635237/answer/756144739

  • 网络设计很重要,外加不要初始化任何用不到的变量,因为 PyTorch 的初始化和 forward 是分开的,他不会因为你不去使用,而不去初始化
  • 合适的num_worker : Pytorch 提速指南 - 云梦的文章 - 知乎 https://zhuanlan.zhihu.com/p/39752167(这里也包含了一些其他细节上的讨论)

模型设计

来自 ShuffleNetV2 的结论:(内存访问消耗时间,memory access cost 缩写为 MAC

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

其他:

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


推理加速

半精度与权重量化

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

  • TensorRT是 NVIDIA 提出的神经网络推理(Inference)引擎,支持训练后 8BIT 量化,它使用基于交叉熵的模型量化算法,通过最小化两个分布的差异程度来实现
  • Pytorch1.3 开始已经支持量化功能,基于 QNNPACK 实现,支持训练后量化,动态量化和量化感知训练等技术
  • 另外 Distiller 是 Intel 基于 Pytorch 开源的模型优化工具,自然也支持 Pytorch 中的量化技术
  • 微软的 NNI 集成了多种量化感知的训练算法,并支持 PyTorch/TensorFlow/MXNet/Caffe2 等多个开源框架

【参考】:


网络 inference 阶段 Conv 层和 BN 层融合

【参考】

时间分析

  • Python 的 cProfile 可以用来分析。(Python 自带了几个性能分析的模块: profilecProfilehotshot,使用方法基本都差不多,无非模块是纯 Python 还是用 C 写的)


项目推荐

  • 基于 Pytorch 实现模型压缩(https://github.com/666DZY666/model-compression):
  • 量化:8/4/2 bits(dorefa)、三值/二值(twn/bnn/xnor-net)
  • 剪枝:正常、规整、针对分组卷积结构的通道剪枝
  • 分组卷积结构
  • 针对特征二值量化的BN融合

PyTorch提速

原始文档:https://www.yuque.com/lart/ugkv9f/ugysgn
声明:大部分内容来自知乎和其他博客的分享,这里只作为一个收集罗列。欢迎给出更多建议。 知乎回答(欢迎点赞哦):
搬运到了github上:https://github.com/lartpang/PyTorchTricks
有喜欢的可以三连下,或者继续补充哦!

* pytorch dataloader数据加载占用了大部分时间,各位大佬都是怎么解决的?


changelog

  • 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日:补充几篇关于模型压缩量化的科普性文章
  • 2020年2月7日: 从文章中摘录了一点注意事项, 补充在了代码层面小节
  • 2020年4月30日:

  • 添加了一个github的文档备份
  • 补充了卷积层和BN层融合的介绍的链接
  • 另外这里说明下,对于之前参考的很多朋友的文章和回答,没有把链接和对应的内容提要关联在一起,估计会导致一些朋友阅读时相关的内容时的提问,无法问到原作者,这里深感抱歉。
  • 调整部分内容,将内容尽量与参考链接相对应

  • 2020年5月18日:发现一个之前的错误:non_blocking=False的建议应该是non_blocking=True.


扩展阅读

PyTorch 版 CenterNet 训练自己的数据集

  • 1
  • 0
  • 1690
收藏
暂无评论
小白学CV
大咖

哈工大 ·

  • 55

    关注
  • 101

    获赞
  • 17

    精选文章
近期动态
  • 工业瑕疵检测
文章专栏
  • 小白学CV的专栏
作者文章
更多
  • 深度图像修复的一个新突破
    532
  • 使用 pytorch 时,训练集数据太多达到上千万张,Dataloader 加载很慢怎么办?
    1.7k
  • 5 个方法让你的模型加速(附代码解析)
    859
  • 基于 OpenCV 的图像强度操作
    615
  • YOLO-v4 目标检测实时手机端实现
    747