0°
PyTorch入门实战教程

PyTorch 源码浅析(四)

自动微分的算符库THNN和THCUNN

在不考虑自动微分的引擎的情况下,实际上想要实现简单的自动微分很简单。只要将不同的算符实现为具有forward方法,和backward方法的类型就可以了,然后用某个引擎来控制调用的顺序(或者说遍历计算图的顺序)。

在这里我先简要地讲一些计算图(以下来自之前在安庆的统计物理workshop的slide,相关图片参考了Cornell的cs5740,但是是我自己画的,应该不用侵删...)

首先一个计算图定义是,一个具有如下性质的有向图:

  • 边:某个函数变量,或者说某个函数的依赖
  • 具有输入边的点:某个函数(或者说算符)
  • 有输出边的点:某个变量

举个例子的话,对于一个简单的表达式 x^T A x 可以表达为如下的计算图

计算图的求值分为前向传播( forward propagation )和后向传播( backward propagation ),分别用来求输出值和梯度。大致的过程就是对叶子节点赋值,然后将节点上函数计算的结果作为值传到下一个节点上,直到整个图的节点都被遍历过。这其中根据图中的结构(比如有圈,loop)可以进行优化,tensorflow等框架就会对这些情况优化。这里略去,下面根据几个图大致看一下这个过程:

(其实这个放slide时候是有动画效果的,anyways,有空再做一个gif)

所以以上的结构使得我们有两种方式去建立一个计算图:

  • 静态图方案:
    • 先定义图的结构,然后给叶子节点赋值(这也是tensorflow中placeholder的由来)
    • 然后根据叶子节点的赋值进行forward
  • 动态图方案,在forward的同时建立图的结构(也就是所谓的动态图)

然后类似forward,可以用后向传播算梯度

以上的过程,使得我们实际上只需要在代码里对每个节点(算符)实现两个方法forward和backward就能够实现前向传播和后向传播。举个THNN里的例子

前面我们已经讲过了,PyTorch的C代码中,下划线前面是类型名称,后面是方法名称,这里updateOutput就是forward方法,而backward方法如下

但是对于有参数的节点例如linear层,convolution层,使用C语言管理参数是很麻烦的,我们希望这里的实现更加干净,所以还需要另外一个方法单独计算参数的梯度,在linear层等有参数的层里就额外有一个accGradParameters方法

这样等到我们将其封装到Python之后,就可以用Python的垃圾回收等功能去管理参数了。THCUNN中的实现基本一致,就是使用的是CUDA,会带__host__, __device__之类的标示符。

源码浅析系列目录

PyTorch 源码浅析(一)

PyTorch 源码浅析(二)

PyTorch 源码浅析(三)

PyTorch 源码浅析(四)

文章来源:罗秀哲知乎专栏

PyTorch入门实战教程
0 条回复 A 作者 M 管理员
    所有的伟大,都源于一个勇敢的开始!
欢迎您,新朋友,感谢参与互动!欢迎您 {{author}},您在本站有{{commentsCount}}条评论