数据pipeline优化

发表于 分类于
本文字数： 1.7k 阅读时长 ≈ 3 分钟

之前训练的时候只是把图像数据保存在磁盘上，通过torchvision Dataloader方式进行批量加载。最近遇到一个问题，就是批量加载ImageNet数据时发现程序运行很慢，但是cpu/gpu的使用率并不高，查询过后发现是遇到了磁盘io的瓶颈

除了使用更好的硬件（比如用固态硬盘替代机械硬盘）来解决问题外，在网上还找到了不少的软件优化方法：

  1. 打开dataloade pin_memory
  2. 替换数据预处理库
  3. 使用lmdb加快磁盘到cpu内存的io
  4. 使用prefetcher加快cpu内存到gpu内存的io
  5. 使用dali加速数据预处理
  6. 其他

data pipeline

通常在gpu上训练深度网络模型，数据从磁盘到gpu的整个过程大体可分为如下步骤：

  1. 从磁盘加载数据到cpu；
  2. 在cpu上进行数据预处理；
  3. 从cpu加载数据到gpu；
  4. 在gpu上进行深度模型训练。

pin_memory

使用torchvision Dataloader时，可以打开pin_memory设置，其目的是将处理后的数据放入到内存锁页区，更方便加载到gpu，从而提高训练速度

当计算机的内存充足的时候，设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False

数据预处理库

目前torchvision默认使用PIL进行图像处理，为了加速数据处理速度，可以使用其他后端进行处理

  1. Pillow。安装后可以直接替代PIL作为torchvision后端
  2. imgaug。提供了丰富的图像处理操作
  3. albumentations-team/albumentations。综合了多个图像处理库（torchvision/imgaug/opencv），实现最快的图像操作

lmdb

最简单的方式就是先把所有训练都读入到cpu内存上，不过这种方式需要大内存，大部分的训练环境都不符合这个条件。lmdb拥有类似的实现方式，首先在磁盘上将所有数据保存为单个文件，然后通过内存映射的方式访问文件，因为文件内寻址的开销非常小，使用指针运算就能实现，从而能够提高数据加载速度

相关实现：

prefetcher

prefetcher是一个数据预取器，通过并行传输的方式，将数据提前从cpu内存传输到gpu内存中，以此来提高训练速度

注意：使用prefetcher时要确保打开了pin_memory设置

相关实现：

dali

dali是nvidia开发的一个加速数据预处理的gpu加速库。如果要使用dali进行加速，那么需要重新构建数据加载流程

相关实现：Using DALI in PyTorch

其他

在调查过程中，发现ImageNet的数据大小不一，并且相对于网络要求的输入大小而言都比较大，所以可以预先进行图像缩放，然后再进行训练，这样也能够有效的提高训练速度

相关阅读