这篇文章主要是之前一段时间的总结，内容是有关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 - 博客园
