PyTorch 提速四倍!提高 DALI 利用率,创建基于 CPU 的 Pipeline

技术讨论 sophie ⋅ 于 1周前 ⋅ 55 阅读

在过去的几年里,深度学习硬件方面取得了巨大的进步,Nvidia的最新产品Tesla V100和Geforce RTX系列包含专用的张量核,用于加速神经网络中常用的操作。

特别值得一提的是,V100有足够的能力以每秒数千张图的速度训练神经网络,这使得基于ImageNet数据集小模型在单GPU上训练只需几小时,与2012年在ImageNet上训练AlexNet模型所花费的5天时间相比简直是天壤之别!

然而,强大的GPU使数据预处理管道不堪重负。为了解决这个问题,Tensorflow发布了一个新的数据加载器:tf.data.Dataset,用C++编写,并使用基于图的方法将多个预处理操作链接在一起。

另一方面,PyTorch使用在PIL库上用Python编写的数据加载器,既方便优灵活,但在速度上有欠缺(尽管PIL-SIMD库确实稍微改善了这种情况)。

进入NVIDIA数据加载器(DALI):旨在消除数据预处理瓶颈,允许训练和推理全速运行。DALI主要用于在GPU上的预处理,但是大多数操作也在CPU上有快速实现。本文主要关注PyTorch,但是DALI也支持Tensorflow、MXNet和TensorRT,尤其是TensorRT有高度支持。它允许训练和推理步骤使用完全相同的预处理代码。需注意,不同的框架(如Tensorflow和PyTorch)通常在数据加载器之间有很小的差异,这可能会影响准确性。

本文是Medium上一位博主展示了一些技术来提高DALI的使用率并创建了一个完全基于CPU的管道。这些技术用于保持长期的内存稳定,并且与DALI包提供的CPU和GPU管道相比,可以增加50%的批处理大小。

DALI长期内存使用

第一个问题是,RAM的使用随着训练时间的增加而增加,这会导致OOM错误(即使是在拥有78GB RAM的VM上),并且尚未修正。

file

唯一解决方案是重新import DALI并每次重构训练和验证通道:

del self.train_loader, self.val_loader, self.train_pipe, self.val_pipe

请注意,使用这种方法,DALI仍然需要大量RAM才能获得最好的结果。考虑到如今RAM的价格,这并不是什么大问题。从下表可以看出,DALI的最大批大小可能比TorchVision低50%:

file

接下来的部分涉及降低GPU占用率的方法。

构建一个完全基于CPU的Pipeline

让我们首先看看示例CPU管道。当不考虑峰值吞吐量时,基于CPU的管道非常有用。CPU训练管道只在CPU上执行解码和调整大小的操作,而CropMirrorNormalize操作则在GPU上运行。由于仅仅是传输输出到GPU与DALI就使用了大量的GPU内存,为了避免这种情况,我们修改了示例CPU管道,使其完全运行在CPU上:

class HybridTrainPipe(Pipeline):

DALI管道现在在CPU上输出一个8位张量。我们需要使用PyTorch来完成CPU-> GPU的传输、浮点数的转换和归一化。这最后两个操作是在GPU上完成的,快速并且减少了CPU -> GPU内存带宽需求。在转移到GPU之前,尝试过固定这个张量,但是没有得到任何性能提升,把它和一个预存器放在一起:

def _preproc_worker(dali_iterator, cuda_stream, fp16, mean, std, output_queue, proc_next_input, done_event, pin_memory):

基于GPU的Pipeline

测试中,在类似最大批处理大小下,上述CPU管道的速度大约是TorchVision数据加载器的两倍。CPU管道可以很好地与像ResNet50这样的大型模型一起工作;然而,当使用像AlexNet或ResNet18这样的小模型时,CPU更好。GPU管道的问题是最大批处理大小减少了近50%,限制了吞吐量。

一种显著减少GPU内存使用的方法是将验证管道与GPU隔离直到最后再调用。这很容易做到,因为我们已经重新导入DALI,并在每个epoch中重新创建数据加载器。

更多小提示

在验证时,将数据集均分的批处理大小效果最好,这避免了在验证数据集结束时还需要进行不完整的批处理。

与Tensorflow和PyTorch数据加载器类似,TorchVision和DALI管道不会产生相同的输出—您将看到验证精度略有不同。我发现这是由于不同的JPEG图像解码器。另一方面,DALI支持TensorRT,允许使用完全相同的预处理来进行训练和推理。

对于峰值吞吐量,尝试将数据加载器的数量设置为number_of_virtual_CPU核心,2个虚拟核对应1个物理核。

如果你想要绝对最好的性能并且不需要有类似TorchVision的输出,尝试关闭DALI三角形插值。

不要忘记磁盘IO。确保您有足够的内存来缓存数据集和/或一个非常快的SSD。DALI读取高达400Mb/s !

合并

为了方便地集成这些修改,我创建了一个data loader类,其中包含这里描述的所有修改,包括DALI和TorchVision后端。使用很简单。实例化数据加载程序:

dataset = Dataset(data_dir,

然后得到训练和验证数据装载器:

train_loader = dataset.get_train_loader()

在每个训练周期结束时重置数据加载器:

dataset.reset()

或者,验证管道可以在模型验证之前在GPU上重新创建:

dataset.prep_for_val()

基准

以下是使用ResNet18的最大批量大小:

file

因此,通过应用这些修改,DALI可以在CPU和GPU模式下使用的最大批处理大小增加了约50%!

这里是一些使用Shufflenet V2 0.5和批量大小512的吞吐量图:

file

这里是一些使用DALI GPU管道训练各种网络,包括在TorchVision:

file

所有测试都在谷歌Cloud V100实例上运行,该实例有12个vcpu(6个物理核心),78GB RAM,并使用Apex FP16培训。要重现这些结果,请使用以下参数:

— fp16 — batch-size 512 — workers 10 — arch “shufflenet_v2_x0_5 or resnet18” — prof — use-dali

所以,DALI使得单核特斯拉V100可以达到接近4000张/秒的图像处理速度!这达到了Nvidia DGX-1的一半多一点(它有8个V100 gpu),尽管我们使用了小模型。对我来说,能够在几个小时内在一个GPU上运行ImageNet是生产力进步。

本文中提供的代码如下:

https://github.com/yaysummeriscoming/DALI\_pytorch\_demo

相关报道:

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

-END-

来源:大数据文摘@微信公众号
编译:赵吉克

*延伸阅读

微信公众号: 极市平台(ID: extrememart )
每天推送最新CV干货

成为第一个点赞的人吧 :bowtie:
回复数量: 0
暂无回复~
您需要登陆以后才能留下评论!