自动微分的算符库THNN和THCUNN
在不考虑自动微分的引擎的情况下，实际上想要实现简单的自动微分很简单。只要将不同的算符实现为具有forward方法，和backward方法的类型就可以了，然后用某个引擎来控制调用的顺序（或者说遍历计算图的顺序）。
在这里我先简要地讲一些计算图（以下来自之前在安庆的统计物理workshop的slide，相关图片参考了Cornell的cs5740，但是是我自己画的，应该不用侵删…）
首先一个计算图定义是，一个具有如下性质的有向图：
- 边：某个函数变量，或者说某个函数的依赖
- 具有输入边的点：某个函数（或者说算符）
- 有输出边的点：某个变量
举个例子的话，对于一个简单的表达式 可以表达为如下的计算图
计算图的求值分为前向传播( forward propagation )和后向传播( backward propagation )，分别用来求输出值和梯度。大致的过程就是对叶子节点赋值，然后将节点上函数计算的结果作为值传到下一个节点上，直到整个图的节点都被遍历过。这其中根据图中的结构（比如有圈,loop）可以进行优化，tensorflow等框架就会对这些情况优化。这里略去，下面根据几个图大致看一下这个过程：
（其实这个放slide时候是有动画效果的，anyways，有空再做一个gif）
所以以上的结构使得我们有两种方式去建立一个计算图：
- 静态图方案：
- 先定义图的结构，然后给叶子节点赋值（这也是tensorflow中placeholder的由来）
- 然后根据叶子节点的赋值进行forward
- 动态图方案，在forward的同时建立图的结构（也就是所谓的动态图）
然后类似forward，可以用后向传播算梯度
以上的过程，使得我们实际上只需要在代码里对每个节点（算符）实现两个方法forward和backward就能够实现前向传播和后向传播。举个THNN里的例子
1 2 3 4 5 6 7 | void THNN_(Sigmoid_updateOutput)( THNNState *state, THTensor *input, THTensor *output) { THTensor_(sigmoid)(output, input); } |
前面我们已经讲过了，PyTorch的C代码中，下划线前面是类型名称，后面是方法名称，这里updateOutput就是forward方法，而backward方法如下
1 2 3 4 5 6 7 8 9 10 11 12 13 | void THNN_(Sigmoid_updateGradInput)( THNNState *state, THTensor *gradOutput, THTensor *gradInput, THTensor *output) { THNN_CHECK_NELEMENT(output, gradOutput); THTensor_(resizeAs)(gradInput, output); TH_TENSOR_APPLY3(real, gradInput, real, gradOutput, real, output, real z = *output_data; *gradInput_data = *gradOutput_data * (1. - z) * z; ); } |
但是对于有参数的节点例如linear层，convolution层，使用C语言管理参数是很麻烦的，我们希望这里的实现更加干净，所以还需要另外一个方法单独计算参数的梯度，在linear层等有参数的层里就额外有一个accGradParameters方法
1 2 3 4 5 6 7 8 9 10 11 | void THNN_(Linear_accGradParameters)( THNNState *state, THTensor *input, THTensor *gradOutput, THTensor *gradInput, THTensor *weight, THTensor *bias, THTensor *gradWeight, THTensor *gradBias, THTensor *addBuffer, accreal scale_)； |
这样等到我们将其封装到Python之后，就可以用Python的垃圾回收等功能去管理参数了。THCUNN中的实现基本一致，就是使用的是CUDA，会带__host__, __device__之类的标示符。
源码浅析系列目录
文章来源：罗秀哲知乎专栏
本站微信群、QQ群(三群号 726282629)：
评论列表(1)