torch.backends.cudnn.benchmark的设置技巧

背景知识

在说 torch.backends.cudnn.benchmark 之前,我们首先简单介绍一下 cuDNN。cuDNN 是英伟达专门为深度神经网络所开发出来的 GPU 加速库,针对卷积、池化等等常见操作做了非常多的底层优化,比一般的 GPU 程序要快很多。大多数主流深度学习框架都支持 cuDNN,PyTorch 自然也不例外。在使用 GPU 的时候,PyTorch 会默认使用 cuDNN 加速。但是,在使用 cuDNN 的时候,torch.backends.cudnn.benchmark 模式是为 False。所以就意味着,我们的程序可能还可以继续提速!

卷积层是卷积神经网络中的最重要的部分,也往往是运算量最大的部分。如果我们可以在底层代码中提升卷积运算的效率的话,就可以在不改变给定的神经网络结构的情况下,大大提升其训练和预测的速度。

对于卷积这个操作来说,其实现方式是多种多样的。最简单的实现方式就是使用多层循环嵌套,对于每张输入图像,对于每个要输出的通道,对于每个输入的通道,选取一个区域,同指定卷积核进行卷积操作,然后逐行滑动,直到整张图像都处理完毕,这个方法一般被称为 direct 法,这个方法虽然简单,但是看到这么多循环,我们就知道效率在一般情况下不会很高了。除此之外,实现卷积层的算法还有基于 GEMM (General Matrix Multiply) 的,基于 FFT 的,基于 Winograd 算法的等等,而且每个算法还有自己的一些变体。在一个开源的 C++ 库 triNNity 中,就实现了接近 80 种的卷积前向传播算法!

每种卷积算法,都有其特有的一些优势,比如有的算法在卷积核大的情况下,速度很快;比如有的算法在某些情况下内存使用比较小。给定一个卷积神经网络(比如 ResNet-101),给定输入图片的尺寸,给定硬件平台,实现这个网络最简单的方法就是对所有卷积层都采用相同的卷积算法(比如 direct 算法),但是这样运行肯定不是最优的;比较好的方法是,我们可以预先进行一些简单的优化测试,在每一个卷积层中选择最适合(最快)它的卷积算法,决定好每层最快的算法之后,我们再运行整个网络,这样效率就会提升不少。

这里有一个问题,为什么我们可以提前选择每层的算法,即使每次我们送入网络训练的图片是不一样的?即每次网络的输入都是变化的,那么我怎么确保提前选出来的最优算法同样也适用于这个输入呢?原因就是,对于给定输入来说,其具体值的大小是不影响卷积的运行时间的,只有其尺寸才会影响。举例来说,我们只要固定输入大小都是 (8, 64, 224, 224),即 batch_size 为 8,输入的通道为 64,宽和高为 224,那么卷积层的运行时间都是几乎不变的,无论其中每个像素具体的值是 0.1 还是 1000.0。

这样的话,因为我们固定了模型输入的尺寸大小,所以对每个卷积层来说,其接受的输入尺寸都是静态的,固定不变的,在提前做优化的时候我们只要使用随机初始化的相应尺寸的输入进行测试和选择就行了。

torch.backends.cudnn.benchmark!

说了这么多背景知识,但和 cudnn.benchmark 有何联系呢?实际上,设置这个 flag 为 True,我们就可以在 PyTorch 中对模型里的卷积层进行预先的优化,也就是在每一个卷积层中测试 cuDNN 提供的所有卷积实现算法,然后选择最快的那个。这样在模型启动的时候,只要额外多花一点点预处理时间,就可以较大幅度地减少训练时间。

这岂不是,用 cudnn.benchmark 一时爽,一直用一直爽吗?其实不然,在某些情况,使用它可能会大大增加运行时间!在背景知识里面我们已经提到过,但是在这里我们更加具体的定义一下,到底哪些因素会影响到卷积层的运行时间。

  1. 首先,当然是卷积层本身的参数,常见的包括卷积核大小,stride,dilation,padding ,输出通道的个数等;
  2. 其次,是输入的相关参数,包括输入的宽和高,输入通道的个数等;
  3. 最后,还有一些其他的因素,比如硬件平台,输入输出精度、布局等等。

我们定义一个卷积场景的参数主要包括 (1) 和 (2),因为在同一个程序中 (3) 往往都是相同的,我们暂且忽略不计。不同的卷积场景有不同的最优卷积算法,需要分别进行测试和选择。

据此我们可以看出来,首先如果我们的网络模型一直变的话,那肯定是不能设置 cudnn.benchmark=True 的。因为网络结构经常变,每次 PyTorch 都会自动来根据新的卷积场景做优化:这次花费了半天选出最合适的算法出来,结果下次你结构又变了,之前就白做优化了。不仅如此,还得要根据这个新的结构继续做选择最高效的算法组合,又花费不少的时间。这样反而会大大降低效率。

另外,我们输入的大小也不能变。对于一个卷积层,这次的输入形状比如是 (8, 3, 224, 224),下次换成了 (8, 3, 112, 112),那也不行。输入的情况变了,最优的算法不一定适用了(比如有的算法在大尺寸输入情况下速度快),PyTorch 还是会重新寻找最优算法的。注意,这里的 batch size,输入通道,图片大小都不能变。

不过一般的 CV 模型来说,网络的结构一般是不会动态变化的,其次,图像一般都 resize 到固定的尺寸,batch size 也是固定的。所以,在大部分情况下,我们都可以在程序中加上这行神奇的代码,来减少运行时间!

等等,这行代码要加在哪里?

之前在网上看到过一些博客提到使用 cudnn.benchmark=True ,但是没有明确说明这段代码要放到哪里。这里我怕有的读者和我当时一样,所以就加了这一小部分。其实一般加在开头就好,比如在设置使用 GPU 的同时,后边补一句:

if args.use_gpu and torch.cuda.is_available():
    device = torch.device('cuda')
    torch.backends.cudnn.benchmark = True
else:
    device = torch.device('cpu')

......
......

当然某些情况下也可以在程序中多次改变 torch.backends.cudnn.benchmark 的值,玩点花样什么的。

原文链接:https://zhuanlan.zhihu.com/p/73711222




如果这篇文章帮助到了你,你可以请作者喝一杯咖啡

原文地址:https://www.cnblogs.com/sddai/p/14846060.html