写点什么

OneFlow 源码解析:自动微分机制

作者:OneFlow
  • 2022-12-02
    重庆
  • 本文字数:11033 字

    阅读完需:约 36 分钟

OneFlow源码解析:自动微分机制

撰文 | 郑建华

更新|赵露阳、王迎港


深度学习框架一般通过自动微分(autograd)机制计算梯度并反向传播。本文尝试通过一个简单的例子,粗浅地观察一下 OneFlow 的 autograd 的实现机制。


1

自动微分基础


自动微分相关的资料比较多,个人感觉自动微分的原理介绍(https://mp.weixin.qq.com/s/BwQxmNoSBEnUlJ1luOwDag)这个系列及其引用的资料对相关背景知识的介绍比较完整清晰。


下面分几种情况对梯度传播的原理做一些直观解释。

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 的梯度,仍可以按链路拆解,一条链路前后节点间的梯度是乘积关系,传入的多条链路梯度是加和关系。

这篇博客(https://blog.paperspace.com/pytorch-101-understanding-graphs-and-automatic-differentiation/)中有一个几乎一样的拓扑图,给出了部分权重参数的梯度公式。


autograd 中 tensor 相关的一些基本概念

2.1 叶子节点


OneFlow 的 autograd 文档(https://docs.oneflow.org/en/master/basics/05_autograd.html)中介绍了 leaf node 和 root node 的概念。只有输出、没有输入的是 leaf node,只有输入、没有输出的是 root node。


个人理解,如果把 weight、bias、data 视为计算图的一部分,这些节点就是叶子节点(op 不是叶子节点)。尤其是从反向计算图的视角(https://discuss.pytorch.org/t/what-is-the-purpose-of-is-leaf/87000/9)看,这些节点的 grad_fn 是空,反向传播到这些节点就会停止。


is_leaf 和 requires_grad 有比较密切的关系,但二者又是独立的。PyTorch 是这样解释的:(https://pytorch.org/docs/stable/generated/torch.Tensor.is_leaf.html#torch.Tensor.is_leaf)


  • 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(https://discuss.pytorch.org/t/what-is-the-purpose-of-is-leaf/87000/9),这个概念只是在阅读代码的时候才会涉及到。


下面的资料可以供进一步参考:


  • What is the purpose of `is_leaf`? (https://discuss.pytorch.org/t/what-is-the-purpose-of-is-leaf/87000)

  • 叶子节点和 tensor 的 requires_grad 参数(https://zhuanlan.zhihu.com/p/85506092


2.2 tensor detach


Tensor 的 detach 方法(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/tensor_impl.cpp#L155)会创建一个新的 tensor,新 tensor 的属性中


  • requires_grad = false

  • is_leaf = true


detach 的意思是从 grad 的反向计算图中把 tensor 分离出来。新的 tensor 与原来的对象共享存储,但不参与反向图的拓扑构造。原有对象的 requires_grad 属性不变。


比如下面的代码,修改一个对象的数据,另一个对象的数据也会改变。

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


示例代码


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

import oneflow as flow# y is scalarx = flow.tensor([-1.0, 2.0], requires_grad=True)y = flow.relu(x).sum()y.backward()print(x.grad)
# y is not scalarx = 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 这一步。


参考资料

  • 自动求梯度 (https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter02_prerequisite/2.3_autograd?id=_233-梯度


  • PyTorch 的 backward 为什么有一个 grad_variables 参数?(https://zhuanlan.zhihu.com/p/29923090


3.1 梯度结果的存储


Tensor 的 grad 属性(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/framework/tensor.cpp#L611),在读取值时调用的是 acc_grad()方法(acc 应该是 accumulate 的缩写)。这样就知道梯度实际存储在哪里,读代码时可以重点关注相关部分。


调用流程如下:


注:图片中的 MirroredTensor 在最新源码中,已经更名为 LocalTensor,其实是一样的。


autograd 相关的类图关系


下图展示了 autograd 相关类的关系


在看 autograd 代码之前,可以参照这个类图,了解其中的结构和关系,有助于理解代码中各个部分的作用。


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


对照上面的类图


站在 tensor 的视角

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

  • 从 y 可以找到反向计算图实际执行梯度计算的类,即 TensorIf -> FunctionNode -> ReLU 这个链路。

  1. FunctionNode 的 backward_fn_包含了 OpExprGradClosure。它只负责计算当前节点的梯度。

  2. 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 成员变量指针。

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

  2. output_meta_data_即 y.autograd_meta_

  3. input_meta_data_即 x.autograd_meta_

  4. 所以 FunctionNode 能获取到上下游的梯度数据并进行读写

  • AutoGradCaptureState 可以存储一些梯度计算需要的状态信息,比如计算 relu 的梯度时需要用到它的前向输出结果 y。

站在反向图执行的视角

  • GraphTask 负责反向图的执行。

  • FunctionNode 只保存必要的数据。

  • GraphTask 基于这些数据,自己构造遍历需要的数据结构,遍历所有节点、执行梯度计算。


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


反向图的执行过程是数据驱动的,数据的存储结构和内容决定了执行的具体动作。


以下讨论只针对 eager 模式。lazy 模式下,反向图的构建是多轮优化 passes 的一部分(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L98)。

之前在讨论Op、Kernel与解释器时已经了解 Interpreter 的作用。只是当时重点关注 op 的执行,忽略了 grad 相关的内容。

GetInterpreter(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter_util.cpp#L67)返回的其实是一个 AutogradInterpreter 对象(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter_util.cpp#L42),在它的 Apply 方法中(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L86),调用内嵌 Interpreter 的同时,也会记录 grad 计算需要的信息。


AutogradInterpreter::Apply 的主要流程如下:



Apply 的第一步会先计算 requires_grad。只要 op 的任一输入的 requires_grad 为 true,op 的输出的 requires_grad 也为 true(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L151-L152)(前提是输出的数据类型支持梯度)。y 的 requires_grad 就是在这里决定的。


比如 y = relu(x),如果数据类型支持梯度,y.requires_grad 就等于 x.requires_grad。


然后会调用内嵌的解释器 internal_执行相关计算。在调用内嵌解释器期间,会临时禁止梯度模式,比如有些 op 可能会嵌套、多次调用解释器(ReluGradFunctor 也会通过解释器执行),这些都不需要梯度逻辑。


需要说明的是,构造 x 时不会执行 grad 相关的逻辑,因为 inputs 的 requires_grad 都是 false,x 的 requires_grad 是在构造的最后才设置的(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/utils/tensor_utils.cpp#L187)。


下面重点看一下几个核心函数的逻辑细节。

5.1 梯度闭包的构建


前面对类图的说明中已经提到,OpExprGradClosure 只负责当前节点的梯度计算。


GetOrCreateOpGradClosure 函数(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_expr.cpp#L146)的核心代码如下:


template<>Maybe<OpExprGradClosure> BuiltinOpExprImpl<UserOpConf>::GetOrCreateOpGradClosure() const {if (!op_grad_func_.get()) {    ...    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(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/common/auto_registration_factory.h#L94)获取预先注册的工厂、创建对象。之前在讨论Op指令在虚拟机中的执行时也看到过类似的注册机制。


这里 op_type_name 的值是 relu,在代码中搜索"relu",可以找到注册 ReLU 的宏(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/gradient_funcs/activation.cpp#L562)。宏展开后的代码如下:


static AutoRegistrationFactory<std::string, OpExprGradFunctionIf>::CreatorRegisterTypeg_registry_var4("relu", ([]() { return new ReLU; }));
复制代码


所以实际返回的对象是 ReLU(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/gradient_funcs/activation.cpp#L200)。其 Init 函数是个空操作。


OpExprGradClosure 只是简单的把 ReLU 存下来供 backward 执行时调用。整个调用流程如下:


5.2 捕获梯度计算需要的数据


调用流程如下:



Capture 函数(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L122)的作用就是为后续的梯度计算保存必要的数据。

需要注意的是,OpExprGradFunction::CaptureIf(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_expr_grad_function.h#L93)中保存的是 detach 的 tensor。这些 tensor 与原来的 tensor 共享数据;可以读写梯度数据,但不会参与反向图的拓扑构造。这个函数把 Interpreter 传过来的 op 的 detached outputs 传给 ReLU::Capture(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_expr_grad_function.h#L128)(就是 relu 的前向输出 y),ReLU::Capture 就把 output[0]存到 ReLUCaptureState 的 saved_tensors_中(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/gradient_funcs/activation.cpp#L209)。因为对于 relu 来说,根据 y 就可以计算梯度。


5.3 保存反向图结构信息


AutogradInterpreter::Apply 中会构造一个 lambada 表达式 backward_fn(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L103-L110),其核心逻辑只有一行 grad_closure->Apply。

这个 lambda 的主要作用就是捕获 grad_closure 这个智能指针。lambda 表达式最终会作为 FunctionNode 的 backward_fn_变量。这样才有类图中 FunctionNode 到 OpExprGradClosure 这条线,才能从 FunctionNode 找到 closue、执行节点的梯度计算。


GetThreadLocalAutogradEngine()->AddNode 这个函数(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L113)很关键,AddNode 的主要任务(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L478)是为 inputs 和 outputs 创建 FunctionNode、并保存反向图遍历需要的数据。其输入参数中的 inputs/outputs,是前向计算的 op 的 inputs/outputs。对于 relu 来说,inputs 就是 x,outputs 就是 y。


在上述示例代码中,对于 x,因为它是叶子节点、也需要梯度,在 AddAccumulateFunctionNode 会将 grad_fn_node 设置为一个空操作的函数(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L508)。之所以是空操作,是因为叶子节点只需要存储梯度、不需要自己计算梯度;它所需要的梯度计算结果会由反向图的上游节点保存到 x.autograd_meta_中。


之后会为 y 构造 GraphFunctionNode 并形成节点连接(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L491)、并保存到 grad_fn_node(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/autograd/autograd_engine.cpp#L495)。需要注意的是,这里的 backward_fn 就是 AutogradInterpreter::Apply 中的 lambda 表达式(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/core/framework/op_interpreter/op_interpreter.cpp#L103-L109)。

需要注意的是,AddBackwardFuncPtr 中的 inputs/outputs 是针对 op 而言,GraphFunctionNode 构造函数中同名变量的是针对 FunctionNode 而言,二者的含义和指向的对象是不一样的。


构造完成后,x 和 y 的 grad_fn_node_字段数据内容如下:


x.grad_fn_node_

name_: accumulate_gradnext_functions_: 空input_meta_data_: 空output_meta_data_: size=1,x.autograd_meta_,requires_grad=true,is_leaf=trueoutput_tensor_infos_: 对应x, relu前向op的inputbackward_fn_: 空函数,AddAccumulateFunctionNode中定义的
复制代码

y.grad_fn_node_

name_: relu_backwardnext_functions_: size=1, x.grad_fn_node, 空操作, AddAccumulateFunctionNode中构造的GraphFunctionNodeinput_meta_data_: x.autograd_meta_, requires_grad=true, is_leaf=trueoutput_meta_data_: size=1, y.autograd_meta_, requires_grad=false, is_leaf=falseoutput_tensor_infos_: 对应y, relu前向op的outputbackward_fn_: AutogradInterpreter::Apply中定义的lambda函数
复制代码


backward 就是根据这些数据,从 roots 出发,完成反向图的遍历。


backward 的入口


在《OneFlow源码阅读4:tensor类型体系与local tensor》(https://segmentfault.com/a/1190000041989895)中提到过,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 flowx = flow.tensor([-1.0, 2.0], requires_grad=True)y = flow.relu(x)y.backward(flow.Tensor([1, 1]))print(x.grad)
复制代码


上述示例代码执行时,Backward(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/autograd/autograd.cpp#L90)的主要参数的值如下:


  • outputs: y, relu 输出的 tensor

  • out_grads: [1, 1]


CheckAndInitOutGrads(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/autograd/autograd.cpp#L49)返回的是 loss 通过当前 op、传到当前节点的梯度。其部分逻辑就是第 3 节讨论的


  • 如果 y 是一个向量,backward 必须传入一个与 y 的 shape 一致的向量(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/autograd/autograd.cpp#L72-L81)。

  • 如果 y 是一个标量,backward 不要参数,框架会自动构造一个全 1 的 tensor(https://github.com/Oneflow-Inc/oneflow/blob/48e511e40e09551408c96722c09bd061ce320687/oneflow/api/python/autograd/autograd.cpp#L70)。


7

autograd.grad


通常,我们都会通过 tensor.backward 或 autograd.backward 触发梯度计算和反向传播,但偶尔也会用到autograd.grad这个接口。autograd.gradautograd.backward很相似,不同之处主要在于:

  • autograd.backward以 outputs(Tensor)作为起点,计算每一个叶子节点的梯度,并且梯度可累积,且保存于对应 inputs(Tensor)的 tensor.grad 上。

  • autograd.grad 接口则是从指定的 outputs 为起点,以指定的 inputs 为终点计算梯度,并按 inputs 参数的顺序返回一个由 inputs 相对应的 grads 构成的 TensorTuple。且梯度是直接获得的,不在 inputs 的 tensor.grad 中累积。


由于autograd.grad就只执行后向计算图中的一部分,在 OneFlow 静态图模式下(lazy mode)TaskGraph 统计入度时就需要做一次剪枝,把不需要计算的结点去掉(参考 TaskGraph::ComputeDependenciesAndPruneNode(oneflow/autograd_engine.cpp at 48e511e40e09551408c96722c09bd061ce320687 · Oneflow-Inc/oneflow

)接口),同时记录每个 inputs 序号,在 FunctionNode::Apply(oneflow/autograd_engine.cpp at 48e511e40e09551408c96722c09bd061ce320687 · Oneflow-Inc/oneflow

)执行后,把需要保存的 grad 及时捕获,最后返回给用户。


8

反向计算中 GraphAutogradEngine 的调用流程


反向图计算的流程分析可以结合 3 类信息

  • 流程代码

  • 上述 x 和 y 的 grad_fn_node_的值

  • 类图以及类之间的关系

RunBackwardAndSaveGrads4LeafTensor函数的几个参数是:

  • outputs: relu 的输出 y

  • out_grads: 用户自己构造的 ones [1, 1]


8.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]
复制代码


8.2 反向图计算任务的构造与执行


FunctionNode 只是记录了反向图的基础信息。RunBackwardAndSaveGrads4LeafTensor 中会再构造一个 GraphTask 对象来表示一次反向计算任务。



8.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]
复制代码

8.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的 Functor 名字是 ReluGrad。对应的 Functor 是ReluGradFunctor(命名空间是 oneflow::one::functional::impl)。


ReluGradFunctor 之后,是基于 Primitive kernel 实现的计算逻辑。ReluGradFunctor 中对应 op 名字是"relu_grad",这个 relu_grad 的注册被包在一个宏定义中,实际上会返回一个 BinaryPrimitiveKernel,这是一种稍显特殊的基于 Primitive 的 kernel,其具体为ep::primitive下的一种BroadcastElementwiseBinary工厂,其对应的 cpu 和 cuda 注册分别位于:

  • oneflow/core/ep/cpu/primitive/broadcast_elementwise_binary.cpp

  • oneflow/core/ep/cuda/primitive/broadcast_elementwise_binary.cu

最终实现位于binary_functor.h(oneflow/binary_functor.h at 48e511e40e09551408c96722c09bd061ce320687 · Oneflow-Inc/oneflow):


template<DeviceType device, typename Src, typename Dst>struct BinaryFunctor<device, BinaryOp::kReluBackwardWithDyY, Src, Dst> {  OF_DEVICE_FUNC BinaryFunctor(Scalar attr0, Scalar attr1) {}
OF_DEVICE_FUNC Dst operator()(Src dy, Src y) const { return static_cast<Dst>((y <= static_cast<Src>(0.0)) ? static_cast<Src>(0.0) : dy); }};
复制代码


至此,完成了梯度计算的逻辑。


8.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.5 临时梯度的释放机制


上述第 5.点中,描述了前向图构建过程中已经存放了对应的 FunctionNode 以及前向 op 所对应的反向backward_fn,实际求梯度、反向传播时,这一个个 backward_fn串联起来构成了反向计算图拓扑,对于其中的每个节点,backward_fn中都可以表示为output_gradsinputs/outputs(可选) -> inputs_grads的一个函数。


其中output_grads 就是链式法则中上游计算的累计梯度,当前节点backward_fn计算完成后,该节点的output_grads就不会再被使用到,从而变成了临时梯度。之后会调用 FunctionNode->ReleaseOutTensorArgs()(oneflow/autograd_engine.cpp at 48e511e40e09551408c96722c09bd061ce320687 · Oneflow-Inc/oneflow)来及时释放该临时梯度。


参考资料

  • https://github.com/Oneflow-Inc/oneflow/tree/48e511e40e09551408c96722c09bd061ce320687

  • OneFlow学习笔记:Autograd解析

  • 自动微分的原理介绍

  • 自动求梯度 (https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter02_prerequisite/2.3_autograd?id=_233-梯度)

  • PyTorch 的 backward 为什么有一个 grad_variables 参数?(https://zhuanlan.zhihu.com/p/29923090)

  • PyTorch 101, Part 1: Understanding Graphs, Automatic Differentiation and Autograd (https://blog.paperspace.com/pytorch-101-understanding-graphs-and-automatic-differentiation/


其他人都在看


欢迎下载体验 OneFlow v0.8.0 最新版本:https://github.com/Oneflow-Inc/oneflow/

用户头像

OneFlow

关注

不至于成为世界上最快的深度学习框架。 2022-03-23 加入

★ OneFlow深度学习框架:github.com/Oneflow-Inc/oneflow ★ OF云平台:oneflow.cloud

评论

发布
暂无评论
OneFlow源码解析:自动微分机制_人工智能_OneFlow_InfoQ写作社区