关于c++:OneFlow源码阅读6自动微分机制

深度学习框架个别通过主动微分（autograd）机制计算梯度并反向流传。本文尝试通过一个简略的例子、浅显地察看一下 OneFlow 的 autograd 的实现机制。
OneFlow 刚刚曾经公布了 v0.8.0。这个文档仍是基于 v0.7.0 的代码。

1 主动微分根底

主动微分相干的材料比拟多，个人感觉主动微分的原理介绍这个系列及其援用的材料对相干背景常识的介绍比拟残缺清晰。
上面分几种状况对梯度流传的原理做一些直观解释。

1.1 stack 网络的梯度流传

以 x -> f -> g -> z 这个 stack 网络为例，依据链式法则：

∂z/∂x = ∂z/∂g * ∂g/∂f * ∂f/∂x

理论运行时，在梯度反向流传过程中：

• z 将 ∂z/∂g 传给 g。
• 如果节点 g 有权重 w 须要计算梯度，就计算∂z/∂w = ∂z/∂g * ∂g/∂w。
• g 须要计算∂g/∂f，再乘以 z 传过来的梯度，将后果传给 f。g 只须要给 f 传递链式乘积的后果，不须要传递各项明细。
• 在训练阶段的前向计算时，g 须要保留 ∂g/∂f 计算依赖的两头后果、以供反向计算时应用。
• 其它节点的流传状况顺次类推。

1.2 简略 graph 的梯度流传

以上面这个简略的 graph 拓扑为例。

在持续之前，须要理解一下多元复合函数微分的根本公式。
下图中，u 和 v 都是对于 x 和 y 的函数，z 是对于 u 和 v 的函数。

依据这个公式能够晓得，z 对 x 的梯度别离沿两条链路流传，z -> u -> x和z -> v -> x，节点 x 将两个梯度之和作为 z 对 x 的梯度。

1.3 简单 graph 的梯度流传

再看一个拓扑略微简单点的例子：

上图能够视为 x -> U -> L，其中 U 是e -> ... -> h 的子图。f -> g的子图能够视为 V。
对于节点 h 来说，它须要把梯度传给 g 和 k。
对节点 e 来说，它须要对 f 和 k 传来的梯度求和，才是 ∂L/∂e。
这样，L 对 x 的梯度，仍能够按链路拆解，一条链路前后节点间的梯度是乘积关系，传入的多条链路梯度是加和关系。

这篇 blog 中有一个简直一样的拓扑图，给出了局部权重参数的梯度公式。

2 autograd 中 tensor 相干的一些基本概念

2.1 叶子节点

OneFlow 的 autograd 文档中介绍了 leaf node 和 root node 的概念。只有输入、没有输出的是 leaf node，只有输出、没有输入的是 root node。
集体了解，如果把 weight、bias、data 视为计算图的一部分，这些节点就是叶子节点（op 不是叶子节点）。尤其是从反向计算图的视角看，这些节点的 grad_fn 是空，反向流传到这些节点就会进行。

is_leaf 和 requires_grad 有比拟亲密的关系，但二者又是独立的。PyTorch 是这样解释的：

• requires_grad=false的节点都是叶子节点。比方 data。
• requires_grad=true的节点如果是用户创立的，也是叶子节点。比方 weight 和 bias。
• 在梯度的反向计算过程中，只有叶子节点的梯度才会被填充。对于非叶子节点，如果要填充梯度信息，须要显式设置retain_grad=true。
• requires_grad=true才会计算、填充梯度。比方y = relu(x)，y 是 op 创立的、不是叶子节点。但如果 x 须要计算梯度，则y.requires_grad==true。但不须要为 y 填充梯度。

对于叶子节点这个概念，目前找到的次要是直观形容，还没看到严格、清晰的定义。也可能是因为用户个别不会间接应用 is_leaf，这个概念只是在浏览代码的时候才会波及到。
上面的材料能够供进一步参考：

• What is the purpose of is_leaf?
• 叶子节点和 tensor 的 requires_grad 参数

2.2 tensor detach

Tensor 的 detach 办法会创立一个新的 tensor。新 tensor 的属性中

• requires_grad = false
• is_leaf = true

detach 的意思是从 grad 的反向计算图中把 tensor 分离出来。新的 tensor 与原来的对象共享存储，但不参加反向图的拓扑结构。原有对象的 requires_grad 属性不变。
比方上面的代码，批改一个对象的数据，另一个对象的数据也会扭转。

import oneflow as flow
y = flow.Tensor([1, 2, 3])
x = y.detach()
x[0] = 4
assert(y[0] == 4)

3 示例代码

本文通过如下代码来察看 OneFlow 的 autograd 机制。

import oneflow as flow
# y is scalar
x = flow.tensor([-1.0, 2.0], requires_grad=True)
y = flow.relu(x).sum()
y.backward()
print(x.grad)

# y is not scalar
x = flow.tensor([-1.0, 2.0], requires_grad=True)
y = flow.relu(x)
y.backward(flow.Tensor([1, 1]))
print(x.grad)

y.backward办法有两种接口：

• 如果 y 是一个标量（比方 loss），不须要传递任何参数。
• 如果 y 是一个向量，须要传入一个与 y 的 shape 统一的向量作为参数。

为什么会有这种区别呢？上面几篇参考资料中对这个问题做了比拟具体的解释。简略的说：

• 如果函数的输入是向量，在反向流传的过程中会造成梯度 tensor shape 的维度收缩，实现简单、性能差。
• 如果函数的输入是标量，反向流传梯度 tensor 的 shape 与参数变量的 shape 统一，不会呈现维度收缩，更容易实现。
• 对于向量版本的 backward，能够假想存在某个 loss 函数，backward 的参数是 loss 流传到 y 这里的梯度。因为前后节点间的梯度是乘积关系，所以用 ones 代替这个假想的梯度，这样计算结果 x.grad 就是 y 对 x 的梯度。

后续将以 y.backward(flow.Tensor([1, 1])) 为例察看一下 autograd 的机制。其反向图只有 x <- y 这一步。

参考资料

• 主动求梯度
• PyTorch 的 backward 为什么有一个 grad_variables 参数？

3.1 梯度后果的存储

Tensor 的 grad 属性，在读取值时调用的是 acc_grad() 办法（acc 应该是 accumulate 的缩写）。这样就晓得梯度理论存储在哪里，读代码时能够重点关注相干局部。
调用流程如下：

4 autograd 相干的类图关系

下图展现了 autograd 相干类的关系

在看 autograd 代码之前，能够参照这个类图，理解其中的构造和关系，有助于了解代码中各个局部的作用。

在 eager 模式下，用户通过 op 的组合逐渐构建出前向计算图。在执行前向计算的过程中，引擎会为 autograd 须要的反向计算图记录必要的信息，在调用 backward 办法时执行这个反向计算图。

对照下面的类图

• 站在 tensor 的视角

  • 前向 op 输入一个 tensor y，即 TensorIf <- ReluFunctor 这部分。

  • 从 y 能够找到反向计算图理论执行梯度计算的类，即 TensorIf -> FunctionNode -> ReLU 这个链路。

    • FunctionNode 的 backward_fn_ 蕴含了 OpExprGradClosure。它只负责计算以后节点的梯度。
    • ReLU 是执行梯度计算的类，它会调用 ReluGradFunctor 这个 op 来执行梯度计算。

• 站在反向图存储的视角

  • 反向图相干的信息在 FunctionNode 中保留。
  • 反向计算图的 root 是 tensor（比方 y 或 loss）的 grad_fn_node_ 变量。
  • FunctionNode 的 next_functions_ 示意反向图的上游节点，以后节点把梯度后果传给这些上游节点。这些 FunctionNode 的连贯就形成了反向图的拓扑构造。
  • tensor 的梯度存储门路是TensorImpl.AutogradMeta.acc_grad_
  • AutogradMeta.current_grad_是反向图上游传递到以后节点的梯度共计。如果 tensor t 输出给 op u 和 v，那么 u 和 v 反传的梯度会累加到current_grad_。current 应该示意截至以后正在计算时的累加和。

  • FunctionNode 尽管并不持有 tensor 实例，但它持有 tensor 的 AutogradMeta 成员变量指针。

    • 基于上述 relu 的例子中的节点 y

      • output_meta_data_即y.autograd_meta_
      • input_meta_data_即x.autograd_meta_
      • 所以 FunctionNode 能获取到上下游的梯度数据并进行读写
  • AutoGradCaptureState 能够存储一些梯度计算须要的状态信息，比方计算 relu 的梯度时须要用到它的前向输入后果 y。

• 站在反向图执行的视角

  • GraphTask 负责反向图的执行。
  • FunctionNode 只保留必要的数据。
  • GraphTask 基于这些数据，本人结构遍历须要的数据结构，遍历所有节点、执行梯度计算。

5 前向计算过程中为 autograd 所做的筹备

反向图的执行过程是数据驱动的，数据的存储构造和内容决定了执行的具体动作。
以下探讨只针对 eager 模式。lazy 模式下，反向图的构建是多轮优化 passes 的一部分。

之前在探讨 Op、Kernel 与解释器时曾经理解 Interpreter 的作用。只是过后重点关注 op 的执行，疏忽了 grad 相干的内容。
GetInterpreter 返回的其实是一个 AutogradInterpreter 对象，在它的 Apply 办法中，调用内嵌 Interpreter 的同时，也会记录 grad 计算须要的信息。

AutogradInterpreter::Apply 的次要流程如下：

Apply 的第一步会先计算 requires_grad。只有 op 的任一输出的 requires_grad 为 true，op 的输入的 requires_grad 也为 true（前提是输入的数据类型反对梯度）。y 的 requires_grad 就是在这里决定的。
比方y = relu(x)，如果数据类型反对梯度，y.requires_grad 就等于 x.requires_grad。

而后会调用内嵌的解释器 internal_ 执行相干计算。在调用内嵌解释器期间，会长期禁止梯度模式，比方有些 op 可能会嵌套、屡次调用解释器（ReluGradFunctor 也会通过解释器执行），这些都不须要梯度逻辑。

须要阐明的是，结构 x 时不会执行 grad 相干的逻辑，因为 inputs 的 requires_grad 都是 false，x 的 requires_grad 是在结构的最初才设置的。

上面重点看一下几个外围函数的逻辑细节。

5.1 梯度闭包的构建

后面对类图的阐明中曾经提到，OpExprGradClosure 只负责以后节点的梯度计算。
GetOrCreateOpGradClosure 函数的外围代码如下：

op_grad_func_.reset(NewObj<std::string, OpExprGradFunctionIf>(proto().op_type_name()));
JUST(op_grad_func_->Init(*this));
return std::make_shared<OpExprGradClosure>(op_grad_func_);

NewObj 会调用 AutoRegistrationFactory 获取事后注册的工厂、创建对象。之前在探讨 Op 指令在虚拟机中的执行时也看到过相似的注册机制。
这里 op_type_name 的值是relu，在代码中搜寻"relu"，能够找到注册 ReLU 的宏。宏开展后的代码如下：

static AutoRegistrationFactory<std::string, OpExprGradFunctionIf>::CreatorRegisterType
    g_registry_var4("relu", ([]() { return new ReLU;}));

所以理论返回的对象是 ReLU。其 Init 函数是个空操作。
OpExprGradClosure 只是简略的把 ReLU 存下来供 backward 执行时调用。

整个调用流程如下：

5.2 捕捉梯度计算须要的数据

调用流程如下：

Capture 函数的作用就是为后续的梯度计算保留必要的数据。
须要留神的是，OpExprGradFunction::CaptureIf 中保留的是 detach 的 tensor。这些 tensor 与原来的 tensor 共享数据；能够读写梯度数据，但不会参加反向图的拓扑结构。这个函数把 Interpreter 传过来的 op 的 detached outputs 传给 ReLU::Capture（就是 relu 的前向输入 y），ReLU::Capture 就 [把 output[0] 存到 ReLUCaptureState 的 saved_tensors_中](https://github.com/Oneflow-In…)。因为对于 relu 来说，依据 y 就能够计算梯度。

5.3 保留反向图构造信息

AutogradInterpreter::Apply 中会结构一个 lambada 表达式 backward_fn，其外围逻辑只有一行 grad_closure->Apply。
这个 lambda 的次要作用就是捕捉 grad_closure 这个智能指针。lambda 表达式最终会作为 FunctionNode 的 backward_fn_ 变量。这样才有类图中 FunctionNode 到 OpExprGradClosure 这条线，能力从 FunctionNode 找到 closue、执行节点的梯度计算。

AddBackwardFuncPtr 这个函数很要害，它的次要工作是为 inputs 和 outputs 创立 FunctionNode、并保留反向图遍历须要的数据。其输出参数中的 inputs/outputs，是前向计算的 op 的 inputs/outputs。
对于 relu 来说，inputs 就是 x，outputs 就是 y。
在上述示例代码中，对于 x，因为它是叶子节点、也须要计算梯度，在 AddAccumulateFunctionNode 会将 grad_fn_node 设置为一个空操作的函数。

之后会为 y 结构 GraphFunctionNode 并造成节点连贯、并保留到 grad_fn_node。须要留神的是，这里的 backward_fn 就是 AutogradInterpreter::Apply 中的 lambda 表达式。
须要留神的是，AddBackwardFuncPtr 中的 inputs/outputs 是针对 op 而言，GraphFunctionNode 构造函数中同名变量的是针对 FunctionNode 而言，二者的含意和指向的对象是不一样的。
结构实现后，x 和 y 的 grad_fn_node_ 字段数据内容如下：

x.grad_fn_node_

op_type_name_: accumulate_grad
next_functions_: 空
input_meta_data_: 空
output_meta_data_: size=1，x.autograd_meta_，requires_grad=true，is_leaf=true
output_tensor_infos_: 对应 x, relu 前向 op 的 input
backward_fn_: 空函数，AddAccumulateFunctionNode 中定义的

y.grad_fn_node_

op_type_name_: relu_backward
next_functions_: size=1, x.grad_fn_node, 空操作, AddAccumulateFunctionNode 中结构的 GraphFunctionNode
input_meta_data_: x.autograd_meta_, requires_grad=true, is_leaf=true
output_meta_data_: size=1, y.autograd_meta_, requires_grad=false, is_leaf=false
output_tensor_infos_: 对应 y, relu 前向 op 的 output
backward_fn_: AutogradInterpreter::Apply 中定义的 lambda 函数

backward 就是依据这些数据，从 roots 登程，实现反向图的遍历。

6 backward 的入口

在 local tensor 的笔记中提到过，Tensor 类在 Python 端通过一层包装，通过 Python 机制为 Tensor 类注册一些办法，backward 就是包装的办法之一。
相干的源代码文件如下

• python/oneflow/framework/tensor.py
• python/oneflow/autograd/__init__.py
• oneflow/python/oneflow/autograd/autograd.py
• oneflow/api/python/autograd/autograd.cpp

C++ 的调用流程如下

这里反复一下本文应用的示例代码

import oneflow as flow
x = flow.tensor([-1.0, 2.0], requires_grad=True)
y = flow.relu(x)
y.backward(flow.Tensor([1, 1]))
print(x.grad)

上述示例代码执行时，Backward 的主要参数的值如下：

• outputs: y, relu 输入的 tensor
• out_grads: [1, 1]

CheckAndInitOutGrads 返回的是 loss 通过以后 op、传到以后节点的梯度。其局部逻辑就是第 3 节探讨的

• 如果 y 是一个向量，backward 必须传入一个与 y 的 shape 统一的向量。
• 如果 y 是一个标量，backward 不要参数，框架会主动结构一个全 1 的 tensor。

7 反向计算中 GraphAutogradEngine 的调用流程

反向图计算的流程剖析能够联合 3 类信息

• 流程代码
• 上述 x 和 y 的 grad_fn_node_ 的值
• 类图以及类之间的关系

RunBackwardAndSaveGrads4LeafTensor 函数的几个参数是：

• outputs: relu 的输入 y
• out_grads: 用户本人结构的 ones [1, 1]

7.1 反向传递过去的梯度的累加

RunBackwardAndSaveGrads4LeafTensor 函数中，PushPartialTensor 的作用就是将 loss 传过来的梯度累加到 autograd_meta_.current_grad_.acc_tensor_。第 4 节中提到，TensorArg.acc_tensor_ 存储的就是 loss 传过来的梯度的共计。这就是 roots（即 y）接管到的梯度，要么是框架主动创立的 ones，要么是用户提供的梯度（通常也是 ones）。

这行代码的逻辑能够用如下伪码示意

outputs[i].impl_.autograd_meta_.current_grad_.acc_tensor_ += out_grads[i]

7.2 反向图计算工作的结构与执行

FunctionNode 只是记录了反向图的根底信息。RunBackwardAndSaveGrads4LeafTensor 中会再结构一个 GraphTask 对象来示意一次反向计算工作。
GraphTask 的构造函数次要是初始化反向图的 roots_ 节点，并将图中各个节点的依赖计数 dependencies_ 置为 0。依据示例代码，roots_就是 y（通常是 loss）。

ComputeDependencies 会对反向图进行深度优先遍历、统计图中各个节点的依赖计数。

GraphTask::Apply 中实现了反向图的遍历逻辑（传入的 save_grad_for_leaf 参数是 true）。当 FunctionNode 的依赖为 0 时，节点才会被放入执行队列，所以对反向图执行的是广度优先遍历，FunctionNode::Apply 执行时，它的依赖都执行结束了。

这个函数中，波及梯度计算逻辑次要包含两局部：

• 调用 node->Apply 执行单个节点的梯度计算
• 调用 node->AccGrad4LeafTensor 存储算好的梯度

7.3 节点的梯度计算

FunctionNode::Apply 中，解决 output_meta_data_的 for 循环的外围逻辑能够用如下伪码示意：

acc_tensor = output_meta_data_[i].current_grad_.acc_tensor_
if (acc_tensor != nullptr) {output_grads[i] = acc_tensor_
} else {output_grads[i] = zeros()}

从中能够看进去，output_grads 的作用就是拷贝上游传过来的梯度数据（指针），作为 backward_fn_的参数。

前面能够看到，backward_fn 的外围逻辑是

// d(y)示意以后节点对 y 的梯度，比方 relu 对其输入 y 的梯度。input_grads = d(y) * output_grads

input_grads 就是以后节点传给上游节点的梯度，调用 backward_fn 时会对它进行赋值。

解决 input_meta_data 的 for 循环的外围逻辑能够用如下伪码示意。本质就是将以后节点传给上游节点的梯度，累加到上游节点的 current_grad 上，从而实现梯度的流传。如果 tensor 输出给多个 op，每个 op 的梯度会加起来。

input_meta_data_[i].current_grad_.acc_tensor_ += input_grads[i]

7.3.1 梯度计算的执行：backward_fn

以下只思考前述示例的 root 节点的执行。也就是 y 对应的 FunctionNode。对于 y 来说，backward_fn 就是 AutogradInterpreter::Apply 中定义的 lambda 表达式。对于 relu 来说，执行过程如下：

之前在 5.1 节曾经确认，OpExprGradClosure::impl_就是 ReLU。
如前所述，backward_fn 的参数中，output_grads 是上游传过来的梯度数据，backward_fn 须要计算 relu 的梯度，二者的乘积赋值给 in_grads。这些参数会始终传递到 ReLU::Apply。

functional::ReluGrad 的 dFunctor 名字是 ReluGrad。对应的 Functor 是 ReluGradFunctor（命名空间是 oneflow::one::functional::impl）。ReluGradFunctor 的 op 名字是 relu_grad，对应 kernel 的注册被包在一个宏定义中，activation_kernels.cpp 中援用了这个宏，并别离针对 float 和 double 进行宏开展。
宏开展后的代码如下

static ::oneflow::user_op::UserOpRegisterTrigger<::oneflow::user_op::OpKernelRegistry>
    g_register_trigger41 =
        ::oneflow::user_op::UserOpRegistryMgr::Get()
            .CheckAndGetOpKernelRegistry("relu_grad")
            .SetCreateFn([]() {
              return user_op::NewOpKernel<BinaryElemwiseXpuKernel<
                  DeviceType::kCPU, ReluGradFunctor<double>, double, double, double>>([](user_op::KernelComputeContext* ctx) {return ReluGradFunctor<double>(); },
                  "dx", "y", "dy");
            })
            .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCPU)
                             && (user_op::HobDataType("dx", 0) == GetDataType<double>::value))
            .SetInplaceProposalFn([](const user_op::InferContext&,
                   const user_op::AddInplaceArgPair& AddInplaceArgPairFn) -> Maybe<void> {for (auto&& maybe___LINE__ = AddInplaceArgPairFn("dx", 0, "dy", 0, true);
                       !maybe___LINE__.IsOk();)
                    return Error(maybe___LINE__.error())
                        .AddStackFrame("oneflow/user/kernels/activation_kernels.cpp", 34,
                                       __FUNCTION__);
                  return Maybe<void>::Ok();});

通过 SetCreateFn 能够晓得，理论执行的 kernel 是 BinaryElemwiseXpuKernel（留神其第二个模版参数是 oneflow::ReluGradFunctor，和 functional::ReluGrad 调用的 Functor 名字雷同、但命名空间不同）。其构造函数中的 FunctorCreateFn 参数就是下面宏开展代码中的

[](user_op::KernelComputeContext* ctx) {return ReluGradFunctor<double>(); }

在它的 Compute 办法中，会将 FunctorCreateFn 执行的后果（也就是 oneflow::ReluGradFunctor）作为参数传给 BinaryElemwiseXpuLauncher 的operator()，所以其中的 functor 就是 oneflow::ReluGradFunctor。这里实现了 relu 的梯度计算逻辑。

7.4 梯度的存储

FunctionNode::Apply 执行结束后，GraphTask::Apply 调用 FunctionNode::AccGrad4LeafTensor 为叶子节点拷贝梯度数据。
在上述例子中，因为 y 不是叶子节点，解决到 y.grad_fn_node_ 时不会进行本质解决。对于 x，会调用 CopyOrAccGrad，这个函数逻辑的伪码模式如下

autograd_meta.acc_grad_ += autograd_meta.current_grad_

autograd_meta.acc_grad_就是 Python 端读到的 x 的梯度。

8 gdb 断点示例

set breakpoint pending on
break oneflow::one::AutogradInterpreter::Apply
break oneflow::one::BuiltinOpExprImpl<oneflow::UserOpConf>::GetOrCreateOpGradClosure
break oneflow::one::OpExprGradFunction<oneflow::one::ReLUCaptureState>::CaptureIf
break oneflow::one::GraphAutogradEngine::AddBackwardFuncPtr
break autograd_engine.cpp:495
break oneflow::one::AddAccumulateFunctionNode
break oneflow::one::GraphFunctionNode::GraphFunctionNode

break oneflow::one::AutogradEngine::RunBackwardAndSaveGrads4LeafTensorIf
break oneflow::one::GraphAutogradEngine::RunBackwardAndSaveGrads4LeafTensor
break oneflow::one::GraphTask::Apply
break oneflow::one::FunctionNode::Apply
break oneflow::one::TensorArg::GetAccTensor
break oneflow::one::functional::impl::BinaryFunctor::operator()
break oneflow::one::functional::TensorProcessor::Apply
break oneflow::one::(anonymous namespace)::CopyOrAccGrad

break elementwise_xpu_kernel.h:48
break oneflow::ReluGradFunctor<float>::operator()

9 参考资料

• oneflow v0.7.0
• OneFlow 学习笔记：Autograd 解析
• OneFlow: AUTOGRAD
• 主动微分的原理介绍
• 主动求梯度
• PyTorch 的 backward 为什么有一个 grad_variables 参数？
• PyTorch 101, Part 1: Understanding Graphs, Automatic Differentiation and Autograd