英伟达DALI加速技巧：让数据预处理速度比原生PyTorch快4倍

6 m' M2 Z& x: P) ^5 I- x V q0 I

作者：Pieterluitjens

0 @+ @( i* r4 U) \* Q- t

参与：一鸣、嘉明、思

你的数据处理影响整个训练速度，如果加上英伟达 DALI 库，处理速度比原生 PyTorch 也能快上四倍。

. q& w% ?# e5 ], M7 t

: d4 o9 F3 x9 f" T2 \) Q( w% ]

打开凤凰新闻，查看更多高清图片

+ Z4 f9 H$ q" r) l3 Y7 l

深度学习的加速上，除了对训练和推理过程的加速之外，还有对数据加载和预处理阶段的加速。这一阶段也尤为关键，因为数据处理 pipeline 的处理速度也影响着整体的流程效率。

近日，有开发者介绍了如何使用 NVIDIA Data Loading Library（DALI），以及使用这个库进行数据预处理加速的方法。DALI 和 TensorFlow 自带的 DataLoader 类似，是一个专门用于加速数据预处理过程的库。

在一篇博客中，开发者表示，该库不仅可以进行 GPU 加速，也可以在 CPU 上构建一个完整的数据预处理流程。如果在 Tesla V100 上做测试，PyTorch+DALI 的处理速度能达到 4000 images/s，比原版 PyTorch 要快近 4 倍。

( J0 N6 u+ N/ h0 i" |' e

支持多个框架，针对预处理

6 r. z3 o, }" u

英伟达数据加载库 DALI 是一个便捷式开源库，用于图像或视频的解码及增强，从而加速深度学习应用。通过并行训练和预处理过程，减少了延迟及训练时间，并为当下流行的深度学习框架中的内置数据加载器及数据迭代器提供了一个嵌入式替代器，便于集成或重定向至不同框架。

" ]& E; b9 a" j9 h0 R _

开发人员可以在 GPU 上运行他们的数据处理通道，从而减少训练神经网络的时间。而 DALI 实现了数据处理 pipeline 可移植，因为可以轻松地重定向至 TensorFlow，PyTorch 和 MXNet。

6 k! J5 f+ e. s; R3 ?5 p. [

0 s4 V l8 W. ?, [2 \; C

DALI 设计之初就是用来帮助用户突破数据预处理的瓶颈，使得模型的训练和推理能够达到最高的效率。其主要设计是用于在 GPU 上的预处理，但大多数操作也可在 CPU 上实现。

主要特征

易于使用的 Python API

在多个 GPU 之间显式地缩放

加快图像分类（ResNet-50）和目标检测分类（SSD）工作流

灵活的计算图使得开发人员可创建自定义的数据处理 pipeline

: d$ ]' n7 s5 B9 g' N* c; H

支持多种格式 - LMDB、RecordIO、TFRecord、COCO、JPEG、H.264 以及 HEVC

3 [% Z4 `+ u2 B' `3 Y

开发人员可添加自定义的图像或视频处理 operator

DALI 在 MXNet 上训练 ResNet50 的性能表现。

使用 DALI 重构数据预处理 pipeline

; l; S: `0 _: y* v

内存占用大是 DALI 在加速数据预处理中面对的一个重要问题。随着批大小的上升，DALI 的内存占用会变得更大。这一问题现在还没有得到修复。

DALI 的内存占用问题。

2 q" O% y" c2 ~4 z

而和其他的数据预处理工具相比，DALI 能够最大处理的批大小相比 TorchVision 少了一半。

因此，使用 DALI 加速数据预处理需要重新构建整个 pipeline。在教程中，作者分别介绍了构建 CPU pipeline 和 GPU pipeline 的方法。

构建 CPU pipepline

0 e* _8 \! |8 ~

使用 CPU 进行数据预处理的时候，常用的做法是让 CPU 处理解码和 resize 的工作，而 CropMirrorNormalize 这样的工作放到 GPU 上。由于使用 DALI 将输出传入 GPU 中就会占用很大的内存，因此作者构建了一个完全基于 CPU 的 pipeline。

: N5 H: ]+ u' }1 O/ a* G; m# ]

首先，重新加载并导入 DALI。

1 R9 P8 B1 H4 ^- Q) z3 a' @

7 ^* M, H( |+ A0 d$ V8 Z! |

然后构建基于 CPU 的 pipeline。

B. L, h; w2 r- q/ D: ~9 ~

在这里设置一下哪些操作应当由 CPU 运行。

, Z: G, V% J. L2 l0 H8 c2 p+ d6 ~

然后开始相应的操作。如解码：

以及 resize：

* N/ e/ }+ d5 a0 q/ _

CropMirrorNormalize 的操作由 GPU 进行。

当然，还需要从 CPU 到 GPU 的转移过程，这一过程使用 PyTorch 进行，作者在教程中也给出了相应的代码。

构建基于 GPU 的 pipeline

当模型较大（如使用 ResNet50）的时候，基于 CPU 的 pipepline 效果较好，而在类似于 AlexNet 或 ResNet18 这样的小模型上，CPU 的处理还是跟不上 GPU 的数据处理速度。其主要原因在于，GPU pipeline 会降低近一半的批大小。

: c7 f; b. Y$ n: \0 W

为了解决这个问题，作者建议将验证 pipepline 从 GPU 中关掉，直到每个 epoch 的最后需要的时候再开启。

2 l- B* H1 c3 [ i+ {' w, `

效果评价

以下是作者在 ResNet18 中使用的最大批量大小：

" T% Z$ C5 T* \

因此通过合理应用这一系列修改，DALI 可以使得最大批量大小在 CPU 或 GPU 模式下提升 50%

而在 Shufflenet V2 0.5 上，如果批大小为 512，则 DALI GPU 和 CPU 能够处理的批大小如下：

: G! m4 ^2 U$ K6 R; S8 C& G6 D0 H

可以看到，即使是 CPU，其处理速度也达到了每秒 1800 张图像，速度超过了 TorchVision。

9 ` h: \' w* k4 F# X6 ]8 I; X9 O$ {

所有的测试都在 Google Cloud V100 实例下运行：包含 12 个 vCPU（6 个物理内核），78GB 内存，以及使用 ApexFP16 进行训练。

+ q- S# ~" @. @5 U7 x# [

因此，通过 DALI，一个 Tesla V100 GPU 就可以达到将近每秒 4000 张图像。接近英伟达有 8 个 V100 GPU 的 DGX-1 的性能，尽管使用的是小模型。

$ `3 s6 m1 b' m

https://towardsdatascience.com/nvidia-dali-speeding-up-pytorch-876c80182440

加入机器之心（全职记者 / 实习生）：hr@jiqizhixin.com

0 Y, x% k7 @/ ~5 l2 @& A

投稿或寻求报道：content @jiqizhixin.com

广告 & 商务合作：bd@jiqizhixin.com

7 z) H" b+ j4 l4 a3 s

4 U a/ ?9 W p) l 2 T0 s- r/ n9 T5 z $ v( q$ Y2 c$ f