参考 torch.backends.cudnn.benchmark ?! - 云+社区 - 腾讯云
大家在训练深度学习模型的时候,经常会使用 GPU 来加速网络的训练。但是说起 torch.backends.cudnn.benchmark
这个 GPU 相关的 flag,可能有人会感到比较陌生。在一般场景下,只要简单地在 PyTorch 程序开头将其值设置为 True
,就可以大大提升卷积神经网络的运行速度。既然如此神奇,为什么 PyTorch 不将其默认设置为 True
?它的适用场景是什么?为什么使用它可以提升效率?答案就在本文之中。
注:因为相关的参考资料比较少,文章的内容是根据我自己的理解和测试的结果总结的,所以如果有错误或者不准确的地方,欢迎大家留言指出。
浓缩版
设置 torch.backends.cudnn.benchmark=True
将会让程序在开始时花费一点额外时间,为整个网络的每个卷积层搜索最适合它的卷积实现算法,进而实现网络的加速。适用场景是网络结构固定(不是动态变化的),网络的输入形状(包括 batch size,图片大小,输入的通道)是不变的,其实也就是一般情况下都比较适用。反之,如果卷积层的设置一直变化,将会导致程序不停地做优化,反而会耗费更多的时间。
行了,对于只想大致了解的读者,可以就此打住了,本文主要说的就是这么一件事,感谢阅读。但对于想刨根问底的读者来说,咱们继续往下看!本文篇幅较长,目录如下:
- 背景知识
- torch.backends.cudnn.benchmark!
- 等等,这行代码要加在哪里?
- PyTorch 中对应的源代码
- 实验结果
- 总结
- 附录:测试代码
背景知识
在说 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。
这样的话,因为我们