撰文｜月踏

更新｜赵露阳

前文《AI杂谈：手推BP》讲了Backward Propagation的数学原理。本文以OneFlow的代码为例，梳理Autograd模块的实现细节。

1

一个求梯度的小例子

先看下面这个简单的例子：

forward pass可以对应到下面的计算图：

 图1

即对应下面公式：

根据前文《AI杂谈：手推‍BP》很容易手动计算出x的梯度值，即： 

x1、x2、x3的计算过程类似，不再赘述，下面看一下OneFlow的执行结果，执行print(x.grad)可得到如下输出：

可以看出，结果和前面公式（3）的计算结果一致，下面通过具体的代码实现来分析OneFlow的Autograd模块。

2

backward接口

上面例子中的python端的backward接口，调用的是python/oneflow/framework/tensor.py中的_backward接口：

可以看到backward只支持eager模式，这是因为graph静态图模式下，计算图是提前编译好的，无需手动通过.backward()调用。flow.autograd.backward()会调用oneflow/api/python/autograd/autograd.cpp中导出的backward方法：

从pybind定义来看，这里面总共导出了两个接口（autograd.backward和autograd.grad）。其中，backward是对所有的requires_grad属性为True的节点求梯度，grad只对指定的叶子结点求梯度，原理上是相同的，本文只以backward为例来看代码的实现，backward接口会调用到同一个文件中的Backward函数：

这里的GetThreadLocalAutogradEngine()可以看作是一个thread_local的单例，位于oneflow/core/autograd/autograd_engine.cpp，返回一个autograd引擎（AutogradEngine）对象的指针：

AutogradEngine是OneFlow的Autograd的核心数据结构，它的继承关系如下：

图2

这里autograd引擎的子类实现有基于栈式的、基于图式的实现，默认使用基于图式的GraphAutogradEngine。从前面代码中可以看到，获取autograd引擎指针后，通过调用RunBackwardAndSaveGrads4LeafTensor函数，位于oneflow/core/autograd/autograd_engine.cpp:L315：

这就真正进入了autograd模块的内部处理流程，后面继续分析。

3

FunctionNode和建立反向图

在进行backward pass时，执行的是一张反向图，反向图中的节点是在forward pass的时候建立的，其中的每个节点被称作FunctionNode，主要数据结构如下：

图3

先说图3中FunctionNode（oneflow/core/autograd/autograd_engine.h:L42），包含next_functions_、input_meta_data_、output_meta_data_这三个数据成员，其中next_functions_表示出边，另外两个表示一些meta信息，下面列几个主要的：

• is_leaf_：是不是叶子节点

• requires_grad_：是不是需要求梯度值

• retain_grad_：对于非叶子节点，是不是保存梯度值

• acc_grad_：在gradient accumulation的的情况下，多个mini-batch的梯度累加

• current_grad_：当前这个batch的梯度值

我们用到的是GraphFunctionNod（oneflow/core/autograd/autograd_engine.cpp:L178）

可见它主要对FunctionNode中的重要数据成员做了初始化，其中input_meta_data_、output_meta_data_中的AutogradMeta信息是从相应的input、output tensor中获取的，tensor通过桥接模式保存了一个TensorImpl对象指针，这个TensorImpl对象则维护了一个AutogradMeta对象。

继续看下FunctionNode中的反向函数backward_fn_，在《OneFlow学习笔记：从Functor到OpExprInterpreter》中讲到了在进行一个op调用的时候会执行AutogradInterpreter::Apply这个函数（oneflow/core/framework/op_interpreter/op_interpreter.cpp:L86），里面会创建这个反向函数：

可以看到反向图节点的名字是以正向图op的type name加上_backward的后缀来组成的，使用AddNode方法来创建FunctionNode（oneflow/core/autograd/autograd_engine.cpp:L356）

可见FunctionNode是挂在Tensor上的，通过Tensor的set_grad_fn_node接口维护到Tensor的数据结构中，在《OneFlow学习笔记：Global View的相关概念和实现》中画过Tensor的继承关系图，FunctionNode就是保存在TensorIf中：

图4

至此，已经理清了FunctionNode中各个成员的作用以及来历，假如以第二节的图1为例来画出对应的反向图的话，如下图所示：

图5

计算好的梯度值会被放到output_meta_data_中得AutogradMeta中，它可以通过tensor的acc_grad、current_grad接口来获取。

反向图的执行流程

接第三节列出的最后一段代码，其中最重要的两句话是：

这里面的graph_task是GraphTask类型，它是一个很重要的数据结构，用来调度反向图中所有FunctionNode的执行，下面列一下它的主要成员：

先看本节开头的graph_task.ComputeDependencies，它主要是在初始化dependencies_这个map，这个map维护了每个FunctionNode的入度信息，再看graph_task.Apply，它主要是在通过拓扑序来访问反向图中的每个FunctionNode，并且对当前的FunctionNode进行各种操作（oneflow/core/autograd/autograd_engine.cpp:L287）

这里最重要的是下面两个语句：

1. node->Apply

2. node->AccGrad4LeafTensor

下面来逐个分析，先看node->Apply（oneflow/core/autograd/autograd_engine.cpp:L143），首先利用output_meta_data_初始化了output_grads，把它作为反向函数的输入，调用反向函数来求梯度值，求出的梯度值暂存在input_grads中，然后再更新到input_meta_data_中：

再看node->AccGrad4LeafTensor，这个函数最终会调用到CopyOrAccGrad，它主要用于在gradient accumulation的时候，多个mini-batch之间把梯度值多累加，和如果有hook函数的的话，使用注册的hook对当前的梯度值进行处理：

（特别感谢同事yinggang中间的各种答疑解惑。本文主要参考代码：https://github.com/Oneflow-Inc/oneflow/commit/a4144f9ecb7e85ad073a810c3359bce7bfeb05e1）

其他人都在看

欢迎下载体验OneFlow v0.7.0：GitHub - Oneflow-Inc/oneflow: OneFlow is a performance-centered and open-source deep learning framework.OneFlow is a performance-centered and open-source deep learning framework. - GitHub - Oneflow-Inc/oneflow: OneFlow is a performance-centered and open-source deep learning framework.https://github.com/Oneflow-Inc/oneflow