您的位置 首页 TMT

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

大数据文摘出品 来源:medium 编译:赵吉克 在过去的几年里,深度学习硬件方面取得了巨大的进步,Nvidia的最新产品TeslaV100和GeforceRTX系列包含专用的张量核,用于加速神经网络中常用的操作。

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

大数据文摘出品

来源:medium

编译:赵吉克

在过去的几年里,深度学习硬件方面取得了巨大的进步,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)通常在数据加载器之间有很小的差异,这可能会影响精确。

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

DALI长期内存使用

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

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

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

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

torch.cuda.synchronize()

torch.cuda.empty_cache()

gc.collect()

importlib.reload(dali)

from dali import HybridTrainPipe,HybridValPipe,DaliIteratorCPU,DaliIteratorGPU

<rebuild DALI pipeline>

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

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

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

构建一个完全基于CPU的管道

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

class HybridTrainPipe(Pipeline):

def __init__(self,batch_size,num_threads,device_id,data_dir,crop,

mean,std,local_rank=0,world_size=1,dali_cpu=False,shuffle=True,fp16=False,

min_crop_size=0.08):

# As we're recreating the Pipeline at every epoch,the seed must be -1 (random seed)

super(HybridTrainPipe,self).__init__(batch_size,seed=-1)

# Enabling read_ahead slowed down processing ~40%

self.input = ops.FileReader(file_root=data_dir,shard_id=local_rank,num_shards=world_size,

random_shuffle=shuffle)

# Let user decide which pipeline works best with the chosen model

if dali_cpu:

decode_device = "cpu"

self.dali_device = "cpu"

self.flip = ops.Flip(device=self.dali_device)

else:

decode_device = "mixed"

self.dali_device = "gpu"

output_dtype = types.FLOAT

if self.dali_device == "gpu" and fp16:

output_dtype = types.FLOAT16

self.cmn = ops.CropMirrorNormalize(device="gpu",

output_dtype=output_dtype,

output_layout=types.NCHW,

crop=(crop,crop),

image_type=types.RGB,

mean=mean,

std=std,)

# To be able to handle all images from full-sized ImageNet,this padding sets the size of the internal nvJPEG buffers without additional reallocations

device_memory_padding = 211025920 if decode_device == 'mixed' else 0

host_memory_padding = 140544512 if decode_device == 'mixed' else 0

self.decode = ops.ImageDecoderRandomCrop(device=decode_device,output_type=types.RGB,

device_memory_padding=device_memory_padding,

host_memory_padding=host_memory_padding,

random_aspect_ratio=[0.8,1.25],

random_area=[min_crop_size,1.0],

num_attempts=100)

# Resize as desired. To match torchvision data loader,use triangular interpolation.

self.res = ops.Resize(device=self.dali_device,resize_x=crop,resize_y=crop,

interp_type=types.INTERP_TRIANGULAR)

self.coin = ops.CoinFlip(probability=0.5)

print('DALI "{0}" variant'.format(self.dali_device))

def define_graph(self):

rng = self.coin()

self.jpegs,self.labels = self.input(name="Reader")

# Combined decode & random crop

images = self.decode(self.jpegs)

# Resize as desired

images = self.res(images)

if self.dali_device == "gpu":

output = self.cmn(images,mirror=rng)

else:

# CPU backend uses torch to apply mean & std

output = self.flip(images,horizontal=rng)

self.labels = self.labels.gpu()

return [output,self.labels]

基于GPU的管道

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

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

更多小提示

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

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

对于前端兆,尝试将数据加载器的数量设置为_virtual_CPU核心的数量,2个虚拟核对应1个物理核。

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

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

合并

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

dataset = Dataset(data_dir,

batch_size,

val_batch_size

workers,

use_dali,

dali_cpu,

fp16)

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

train_loader = dataset.get_train_loader()

val_loader = dataset.get_val_loader()

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

dataset.reset()

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

dataset.prep_for_val()

基准

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

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

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

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

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

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

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

所有测试都在谷歌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返回搜狐,查看更多

责任编辑:

作者: 网站小编

这家伙很懒什么都没说!

为您推荐

返回顶部