这篇文章主要是之前一段时间的总结,内容是有关PyTorch中卷积部分的源码。文章不会很透彻的去研究源码,只是大概地总结一下,主要内容有:
- PyTorch-拓展模块
- PyTorch对于卷积的内部实现
- 为什么有了cudnn还需要PyTorch实现卷积?
很感谢网上的优质博客,正是因为有了知识的共享,人们的生活质量才会不断提高~
本人参考源码实现的卷积链接: [点我跳转],为PyTorch extension做一点小贡献~
1.事件的开始
之前和朋友做一个idea,苦于PyTorch上没有给出想要的轮子。在网上搜寻了一番发现PyTorch有个拓展模块叫C++ extension,即借助PyTorch给的接口,以C++的形式去实现自己想要的功能,甚至可以让你自己编写基于CUDA编程的操作(具体见参考链接4或参考链接5)。
现在有挺多优质的项目在用着拓展模块(例如参考链接7),主要原因是自己造的轮子更适用。本人要做的工作和卷积有关,于是想着能不能借助拓展模块先实现个卷积出来。
但是拓展模块中使用的C++接口说明不太完备,而且网上的有关这个的介绍也不算太多,更多的时候需要观察、自己去总结得出经验。两眼一抹黑,不妨先看看PyTorch是如何实现卷积,然后依葫芦画瓢岂不美哉?
2.PyTorch卷积实现
PyTorch内部有三大块:C10、ATen、Torch,我们主要关注后面两块(点我跳转)。Torch是一个基于C实现的开源项目,在PyTorch中分成以下几个部分:
- TH = TorcH
- THC = TorcH Cuda
- THCUNN = TorcH CUda Neural Network
- THNN = TorcH Neural Network
可以看出,每个模块都有cpu和gpu版,对于神经网络(nn)还有特定的模块(THCUNN,THNN),我们要研究的卷积源码就在其中。
点开THCUNN,看起来眼花缭乱,大致如下
很奇怪的是generic是干嘛用的?为什么里面实现的好像和外面的是一个方法?原因在于Torch本身是一个基于C实现的库,但是C语言既没有面向对象也没有泛型,编写起来难度大很多。但是Torch借助了宏以及命名规则进行巧妙地“伪装”,具体原理这里就不展开了,可以阅读参考链接10,作者写的非常非常详细,而且还附带了个例子。
PyTorch实现的卷积在generic下的SpatialConvolutionMM.cu,我们关注其中三个方法:
- void THNN_(SpatialConvolutionMM_updateOutput) 前向传输
- void THNN_(SpatialConvolutionMM_updateGradInput) 获得对输入的梯度
- void THNN_(SpatialConvolutionMM_accGradParameters) 获得对参数的更新梯度
只要我们搞懂了这三个函数,我们也能借助PyTorch的接口很轻易地实现卷积!
2.1 前向传输
大家直观上理解卷积通常都是用滑窗的形式,但是这样去实现显然很不高效。PyTorch或者Caffe中卷积的实现都是基于一个im2col算法,具体来说,将特征图中每个待卷积的窗口展开成一列并拼接到一起,这样卷积的操作就可以用一个矩阵乘法来代替,如下为一个例子(图来自参考链接11):
为了能更好地理清关系,给出某些重要变量和对应的维:
- 对输入图像的两个方向做padding:padH, padW
- 卷积核在两个方向上移动的步长:dH, dW
- 卷积核权重weights的维度:nOutputPlane * nInputPlane * kH * kW
- 卷积核偏置bias的维度:nOutputPlane
- 输入input的维度:batch_size * nInputPlane * inputHeight * inputWidth
- 输出output的维度:batch_size * nOutputPlane * outputHeight * outputWidth
在已知了输入的维度、卷积核、padding、步长等参数后,可用下面公式计算得出:
接着为了做矩阵乘法,我们先将卷积核的权重reshape成一个矩阵,即:
具体来说,我们都知道有nOutputPlane个卷积核,最后卷积后的特征维度为nOutputPlane。这一步reshape相当于将每个卷积核展平成一个个行向量并拼接在一起构成矩阵。
接着对输入采用im2col算法,该算法会将每一个待卷积的区域展平拼接起来,算法的使用还需要给出padding和步长等参数,原因在于要计算输出维度。算法在具体使用过程中,是以一个循环的形式对batch中的每一个输入张量使用。经过了im2col后,输入维度变化为:
形成了一个矩阵,相当于是将每一块待被卷积的区域拉平成一个列向量,接着拼接在一起,可观察上面的图例来理解。两个矩阵更直观地有:
在完成了矩阵乘法后,reshape成输出的维度即可:
这便是PyTorch中卷积操作的实现。我照着复现并且忽略一些无关紧要的东西,代码大致如下:
为了能和源码同步易于理解,有些临时变量也声明了出来,例如columns张量用于存储im2col的结果,ones张量用于将偏置放入结果矩阵中。代码大部调用了PyTorch拓展模块的api,这些api的声明大部分在[点我跳转]这里找到,不过没给函数的介绍和使用例子。
2.2获得对输入的梯度
神经网络的调参依据的是链式法则,故需要知道损失函数对上一层的求导。要计算对输入的梯度也很简单,需要对卷积结果的梯度以及卷积核权重,这里先给出一些重要定义:
- 输出的梯度gradOutput的维度:batch_size * nOutputPlane * outputHeight * outputWidth
- 输入的梯度gradInput的维度:batch_size * nInputPlane * inputHeight * inputWidth
- 临时变量gradColumns的维度:(nInputPlane×kH×kW) * (outputHeight×outputWidth)
首先明确,无论是对输出还是输入的梯度,其维度和输入input或输出output的维度是一致的。
想了解有关卷积神经网络的反向传播公式推导可以阅读参考链接12。在具体实现的过程会用到col2im算法,没错,就是im2col的逆过程,如下图所示(图来自参考链接1):
首先对卷积核权重做reshape并转置:
接着是对gradOutput做reshape,这里是对batch的每一个output张量做操作:
接下来就是矩阵相乘!再用col2im算法去处理结果,即大功告成:
这就是PyTorch的计算方式,这里借助拓展模块的接口加以实现,大致如下:
2.3获得对权重、偏置的梯度
最后是求对权重和偏置的梯度,这一步完成了整个卷积的流程就打通了。需要的参数有输入input、对卷积输出的梯度gradOutput、卷积核权重weights以及padding和步长。先给出重要参数的维度:
- 卷积核权重的梯度gradWeight的维度:nOutputPlane * nInputPlane * kH * kW
- 卷积核偏置的梯度gradBias的维度:nOutputPlane
根据链式法则,对权重的偏导计算得到就是输入乘以卷积输出。具体运算还需要经过im2col以及一些reshape操作,这里直接给出流程:
---------------------------------------------------------------------------------
对batch中的某一个gradOutput_n=gradOutput[i]有
- 将gradOutput_n reshape成(nOutputPlane, outputHeight * outputWidth)
- 对input采用im2col算法并转置,维度为(outHeight * outWidth, nInputPlane * kW * kH)
- 二者做矩阵乘法并reshape后即完成当前迭代计算(nOutputPlane, nInputPlane, kW, kH),接着同理对batch内所有的梯度累加即可得到gradWeights
---------------------------------------------------------------------------------
偏置的梯度计算就不多说了,大体是gradOutput_n与全1矩阵做乘法即可。代码大致如下:
2.4总结
本节主要描述了下PyTorch源码中对卷积的实现,并以拓展模块的接口复现了一下。代码已经开源在GitHub上[点我跳转],不算是高难度的项目,只是为PyTorch的拓展模块做一点贡献,供学习使用。
3.为什么有了cudnn还需要PyTorch实现卷积?
我复现了以后将它包装成卷积类,尝试了一下效率,慢的令人发指,原因也很简单,上述的实现显然过于简单,而且是对一个batch里的样本串行计算,那么PyTorch的卷积实现在哪?
答案是在ATen文件夹中,我们进入ATen/naive文件夹:
可以看出,里面大量的常见函数,例如MaxPooling、im2col等,当然包括了卷积Convolution。我们打开ATen/naive/Convolution.cpp,找到卷积函数:
按照这样的查找思路,可以发现最后的卷积操作都汇集在了“_convolution”方法中。
该函数在卷积中做了很多判断,来决定使用哪种卷积(其中miopen是AMD开发的gpu加速库,mkldnn是英特尔开发的cpu加速库):
- 是否是depthwise卷积,如果是:
- 是否可用cudnn,是则调用at::cudnn_convolution
- 是否可用miopen,是则调用at::miopen_depthwise_convolution
- 调用at::thnn_conv_depthwise2d
- 是否可用cudnn,如果是:
- 是否是反卷积,是则调用at::cudnn_convolution_transpose
- 调用at::cudnn_convolution
- 是否可用miopen,如果是:
- 是否是反卷积,是则调用at::miopen_convolution_transpose
- 调用at::miopen_convolution
- 是否可用mkldnn,如果是:
- 调用at::mkldnn_convolution
- 调用at::_convolution_nogroup
可以看出,PyTorch支持很多形式的卷积加速,这些个方法都封装在ATen/naive相应的文件夹下,例如有关cudnn的操作就在ATen/naive/cudnn中,想要研究更多有关PyTorch如何调用cudnn可以去该文件夹中瞅瞅:
而这些个方法都会被ATen/naive下的native_functions.yaml标记,在该文件中也能查到api。我们回到问题,那当用户的环境中没有cudnn、miopen、mkldnn时,at::_convolution_nogroup方法用的是什么呢?
在一番研究后发现,at::_convolution_nogroup也是一顿判断,使用了大概这么几个函数:
- at::slow_conv_transpose2d
- at::slow_conv_transpose3d
- at::slow_conv_dilated2d
- at::_nnpack_spatial_convolution
- at::thnn_conv2d
- at::slow_conv_dilated3d
- at::thnn_conv3d
取名就挺好笑的,叫slow。接着我在native_functions.yaml里瞅瞅,看到了官方这么一段话:
原来上一节研究的那些代码就是在这时被使用的,官方承认这种实现是"this is very memory inefficient!"。接着thnn_和slow两种函数的区别在于前者是比较老派的写法,即上一节我们研究的以C实现的卷积,而后者是用C++实现的。
thnn_conv2d等系列函数在native_functions.yaml中如下:
可以看到都是在名为"nn"的python_module中,然后我的好奇心又起来了,这个"nn"究竟在哪?后面在ATen/nn.yaml里找到:
呀SpatialConvolutionMM方法就是我们在上一节研究的文件,看来破案了(应该吧?)。
4.总结
总结一下,PyTorch实现的卷积位于THCUNN/generic/SpatialConvolutionMM.cu,接着PyTorch对外的卷积接口是在ATen/native/cudnn/Conv.cpp中,在这接口里,它判断该使用的卷积加速库(cudnn,mkldnn等),但是如果环境中没有这些加速库,那么就调用PyTorch自己实现的卷积。
PyTorch自己实现的卷积(SpatialConvolutionMM等)在ATen/nn.yaml文件中被“粘贴”(个人理解),然后在ATen/native/native_functions.yaml被“引用”,native_functions.yaml中还包括了各种加速库的调用。有兴趣去研究PyTorch怎么调用cudnn等加速库可以到ATen/native下找对应的文件。
总的来说这是我个人的一个有趣的经历,研究源码的过程学到了不少新知识。应该会有大神会对此不屑一顾,认为这些早都知道了(瑟瑟发抖),希望别被喷吧。以上的内容会包括自己的一些瞎猜,路过的朋友欢迎指正~
参考链接
- PyTorch源码浅析(3):NN
- pytorch/pytorch
- Custom C++ and CUDA Extensions
- Custom C++ and CUDA Extensions
- https://pytorch.org/cppdocs/
- zhanghang1989/PyTorch-Encoding
- A quick tour of Torch internals
- Gemfield:PyTorch ATen代码的动态生成
- 罗秀哲:PyTorch源码浅析(一)
- Making faster · Artificial Inteligence
- 卷积神经网络(CNN)反向传播算法 - 刘建平Pinard - 博客园