机器学习
深度学习（Deep Learning）
Torch (深度学习框架)
PyTorch
PyTorch深度学习（书籍）

使用pytorch时，训练集数据太多达到上千万张，Dataloader加载很慢怎么办?

训练集全是16x16，32x32之类的小图，达到上千万张，训练时发现数据加载很慢很慢很慢！！！看了下CPU 内存 GPU使用情况，发现CPU使用率都跑…
关注者
1,659
被浏览
1,437,257
登录后你可以
不限量看优质回答私信答主深度交流精彩内容一键收藏
查看全部 48 个回答

PyTorch提速

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

changelog

  • 2019年11月29日：更新一些模型设计技巧和推理加速的内容，补充了下apex的一个介绍链接， 另外删了tfrecord，pytorch能用么？这个我记得是不能，所以删掉了（表示删掉:<）
  • 2019年11月30日：补充MAC的含义，补充ShuffleNetV2的论文链接
  • 2019年12月02日：之前说的pytorch不能用tfrecord，今天看到zhihu.com/question/3586下的一个回答，涨姿势了
  • 2019年12月23日：补充几篇关于模型压缩量化的科普性文章
  • 2020年2月7日: 从文章中摘录了一点注意事项, 补充在了代码层面小节
  • 2020年4月30日：
    • 添加了一个github的文档备份
    • 补充了卷积层和BN层融合的介绍的链接
    • 另外这里说明下，对于之前参考的很多朋友的文章和回答，没有把链接和对应的内容提要关联在一起，估计会导致一些朋友阅读时相关的内容时的提问，无法问到原作者，这里深感抱歉。
    • 调整部分内容，将内容尽量与参考链接相对应
  • 2020年5月18日：发现一个之前的错误:non_blocking=False的建议应该是non_blocking=True.
  • 2021年01月06日：调整下关于读取图片数据的一些介绍.
  • 2021年01月13日：补充了一条推理加速的策略. 我觉得我应该先更新github的文档，知乎答案的更新有点麻烦，也没法比较更改信息，就很费劲。

预处理提速

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

IO提速

使用更快的图片处理

  • opencv 一般要比 PIL 要快（但是要注意，PIL的惰性加载的策略使得其看上去open要比opencv的imread要快，但是实际上那并没有完全加载数据，可以对open返回的对象调用其load()方法，从而手动加载数据，这时的速度才是合理的）
  • 对于 jpeg 读取，可以尝试 jpeg4py
  • 存 bmp 图（降低解码时间）
  • 关于不同图像处理库速度的讨论建议关注下这个：

小图拼起来存放（降低读取次数）

对于大规模的小文件读取，建议转成单独的文件，可以选择的格式可以考虑：TFRecord（Tensorflow）、recordIO（recordIO）、hdf5、 pth、n5、lmdb 等等（github.com/Lyken17/Effi）

预读取数据

  • 预读取下一次迭代需要的数据
  • Pytorch1.7中的DataLoader提供了一个参数prefetch_factor ，可以试试
    • Number of sample loaded in advance by each worker. 2 means there will be a total of 2 * num_workers samples prefetched across all workers.

【参考】

借助内存

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

【参考】

借助固态

  • 把读取速度慢的机械硬盘换成 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
  • 网络设计很重要，外加不要初始化任何用不到的变量，因为 PyTorch 的初始化和 forward 是分开的，他不会因为你不去使用，而不去初始化
  • 合适的num_worker ： Pytorch 提速指南 - 云梦的文章 - 知乎 zhuanlan.zhihu.com/p/39（这里也包含了一些其他细节上的讨论）

模型设计

来自 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 自带了几个性能分析的模块： profile 、cProfile 和 hotshot，使用方法基本都差不多，无非模块是纯 Python 还是用 C 写的）

项目推荐

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

扩展阅读

查看全部 48 个回答