Press "Enter" to skip to content

OneFlow学习笔记:Autograd解析

 

一、简介

 

前文《AI杂谈:手推BP》讲了backward propagation的数学原理,本文继续来看Autograd,Autograd可以说是训练框架的核心模块,主流的训练框架都会有这个模块,本文以OneFlow的代码为例,分析梳理一下这个模块的实现细节,特别感谢同事yinggang中间的各种答疑解惑

 

二、一个求梯度的小例子

 

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

 

import oneflow as of
x = of.randn(2, 2, requires_grad=True)
y = x + 100
z = y.sum()
z.backward()

 

这段代码很简单,forward pass可以对应到下面的计算图:

 

 

 

图1

 

即对应下面公式:

 

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

 

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

 

tensor([[1., 1.],
        [1., 1.]], dtype=oneflow.float32)

 

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

 

三、backward接口

 

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

 

def _backward(self, gradient=None, retain_graph=False, create_graph=False):
    if not lazy_mode.is_enabled():
        flow.autograd.backward(self, gradient, retain_graph, create_graph)
    else:
        ...

 

可以看到backward只支持eager模式,会继续调用python/oneflow/autograd/autograd.py中的backward接口:

 

from oneflow._oneflow_internal.autograd import backward as backward_api


def backward(
    outputs: Union[Tensor, Sequence[Tensor]],
    out_grads: Union[Tensor, Sequence[Tensor], None],
    retain_graph: bool = False,
    create_graph: bool = False,
) -> None:
    backward_api(
        convert_to_tensor_tuple(outputs),
        convert_to_tensor_tuple(out_grads),
        retain_graph,
        create_graph,
    )

 

从这里的import语句可以看出,backward_api其实是backward的一个alias,是pybind
导出的一个接口
,定义位于
oneflow/api/python/autograd/autograd.cpp

 

ONEFLOW_API_PYBIND11_MODULE("autograd", m) {
  m.def("backward", &BackwardOrThrow);
  m.def("grad", &GradOrThrow);
}

 

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

 

Maybe<one::TensorTuple> Backward(const one::TensorTuple& outputs, const one::TensorTuple& out_grads,
                                 bool retain_graph, bool create_graph) {
  if (create_graph) { retain_graph = true; }
  auto gradients = CheckAndInitOutGrads(outputs, out_grads);
  one::GetThreadLocalAutogradEngine()->RunBackwardAndSaveGrads4LeafTensorIf(
      outputs, *gradients, retain_graph, create_graph);
  return std::make_shared<one::TensorTuple>(0);
}

 

这里的GetThreadLocalAutogradEngine可以看作是一个thread_local的单例,位于oneflow/core/autograd/autograd_engine.cpp,返回一个GraphAutogradEngine的对象指针:

 

AutogradEngine* GetThreadLocalAutogradEngine() {
  thread_local static GraphAutogradEngine autograd_engine;
  return &autograd_engine;
}

 

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

 

 

 

图2

 

从前面代码看到,调用了GraphAutogradEngine中的RunBackwardAndSaveGrads4LeafTensor函数,位于oneflow/core/autograd/autograd_engine.cpp:

 

Maybe<void> GraphAutogradEngine::RunBackwardAndSaveGrads4LeafTensor(
      const TensorTuple& outputs,
      const TensorTuple& out_grads,
      bool retain_graph,
      bool create_graph) {
  for (int i = 0; i < outputs.size(); ++i) {
    JUST(JUST(outputs.at(i)->current_grad())->PushPartialTensor(out_grads.at(i)));
  }
  GraphTask graph_task(outputs, retain_graph, create_graph);
  graph_task.ComputeDependencies();
  graph_task.Apply(true);
  return Maybe<void>::Ok();
}

 

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

 

四、FunctionNode和建立反向图

 

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

 

 

 

图3

 

先说图3中FunctionNode的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的梯度值

 

我们用到的是GraphFunctionNode,它的构造函数位于oneflow/core/autograd/autograd_engine.cpp:

 

GraphFunctionNode::GraphFunctionNode(
      const std::string& op_type_name,
      function<void(const TensorTuple&, TensorTuple*, bool)>& backward_fn,
      const TensorTuple& inputs, 
      const TensorTuple& outputs) : FunctionNode(op_type_name) {
  input_meta_data_.resize(inputs.size());
  next_functions_->reserve(inputs.size());
  for (int i = 0; i < inputs.size(); ++i) {
    if (inputs.at(i)->requires_grad()) {
      input_meta_data_.at(i) = inputs.at(i)->mut_autograd_meta();
      next_functions_->emplace_back(inputs.at(i)->mut_grad_fn_node());
    }
  }


  output_meta_data_.resize(outputs.size());
  output_tensor_infos_.reserve(outputs.size());
  for (int i = 0; i < outputs.size(); ++i) {
    const auto& autograd_meta =
        NewAutogradMeta(outputs.at(i)->requires_grad(), outputs.at(i)->is_leaf());
    outputs.at(i)->set_autograd_meta(autograd_meta);
    output_meta_data_.at(i) = outputs.at(i)->mut_autograd_meta();
    output_tensor_infos_.emplace_back(TensorInfo(*outputs.at(i)));
  }


  backward_fn_ = backward_fn;
}

 

可见他主要对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,里面会创建这个反向函数:

 

Maybe<void> AutogradInterpreter::Apply(
        const OpExpr& op_expr, 
        const TensorTuple& inputs,
        TensorTuple* outputs, 
        const OpExprInterpContext& ctx) const {
  ...
  autograd::AutoGradMode mode(false);
  JUST(internal_->Apply(op_expr, inputs, outputs, ctx));


  if (requires_grad && !LazyMode::is_enabled()) {
    auto& grad_closure = op_expr.GetOrCreateOpGradClosure();
    JUST(grad_closure->Capture(inputs, *outputs, ctx));


    auto backward_fn = [=](const TensorTuple& out_grads, 
                           TensorTuple* in_grads,
                           bool create_graph) {
        grad_closure->Apply(out_grads, in_grads);
        return Maybe<void>::Ok();
    };
    GetThreadLocalAutogradEngine()->AddBackwardFuncPtr(
      op_expr.op_type_name() + "_backward", backward_fn, inputs, outputs);
  }
  ...
  return Maybe<void>::Ok();
}

 

可以看到反向图节点的名字是以正向图op的type name加上_backward的后缀来组成的,使用AddBackwardFuncPtr来创建FunctionNode,位于oneflow/core/autograd/autograd_engine.cpp:

 

Maybe<FunctionNode> GraphAutogradEngine::AddBackwardFuncPtr(
    string op_type_name,
    function<void(const TensorTuple&, TensorTuple*, bool)>& backward_fn,
    const TensorTuple& inputs, 
    TensorTuple* outputs) {
  // Firstly push function_node of tensor in stack which is leaf and requires_grad
  for (const auto& in_tensor : inputs) {
    if (in_tensor->is_leaf() && in_tensor->requires_grad()) {
      if (!in_tensor->grad_fn_node()) 
        AddAccumulateFunctionNode(in_tensor);
    }
  }


  auto func_node = make_shared<GraphFunctionNode>(
      op_type_name, backward_fn, inputs, *outputs);
  for (const auto& out_tensor : *outputs) {
    out_tensor->set_grad_fn_node(func_node);
  }
  return func_node;
}

 

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

 

 

 

图4

 

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

 

 

 

图5

 

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

 

五、反向图的执行流程

 

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

 

...
graph_task.ComputeDependencies();
graph_task.Apply(true);
...

 

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

 

class GraphTask final {
  bool retain_graph_;
  bool create_graph_;
  std::vector<FunctionNode*> roots_;
  HashMap<FunctionNode*, int> dependencies_;
  HashSet<FunctionNode*> need_execute_;
};

 

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

 

Maybe<void> GraphTask::Apply(bool save_grad_for_leaf) {
  std::queue<FunctionNode*> queue;
  for (auto* node : roots_) {
    if (dependencies_[node] == 0) { queue.push(node); }
  }


  while (!queue.empty()) {
    FunctionNode* node = queue.front();
    queue.pop();
    if (!need_execute_.empty() && need_execute_.find(node) == need_execute_.end()) {
      node->ReleaseOutTensorArgs();
      continue;
    }
    if (!(node->Apply(create_graph_))) { continue; }
    if (save_grad_for_leaf)
        node->AccGrad4LeafTensor(create_graph_);
    node->AccGrad4RetainGradTensor();
    node->ReleaseOutTensorArgs();
    if (!retain_graph_) { node->ReleaseData(); }


    for (const auto& next_grad_fn : *(node->GetNextFunctions())) {
      auto* next_node = next_grad_fn.get();
      dependencies_[next_node] -= 1;
      if (dependencies_[next_node] == 0) { queue.push(next_node); }
    }
  }
  return Maybe<void>::Ok();
}

 

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

 

 

node->Apply

 

node->AccGrad4LeafTensor

 

 

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

 

Maybe<bool> FunctionNode::Apply(bool create_graph) {
  ...
  (*backward_fn_)(output_grads, &input_grads, create_graph);
  for (int i = 0; i < input_meta_data_.size(); ++i) {
    if (input_grads.at(i)) {
      input_meta_data_.at(i)->current_grad()->PushPartialTensor(input_grads.at(i));
    }
  }
  return true;
}

 

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

 

Maybe<void> CopyOrAccGrad(AutogradMeta* autograd_meta, bool autograd_mode) {
  auto current_grad = autograd_meta->current_grad()->GetAccTensor({});


  if (autograd_meta->acc_grad()) {
    const auto& output = functional::Add(
        autograd_meta->acc_grad(), current_grad, 1,
        autograd_meta->is_grad_acc_inplace());
    autograd_meta->set_acc_grad(output);
  } else {
    autograd_meta->set_acc_grad(current_grad);
  }
  for (const auto& hook : autograd_meta->post_grad_accumulation_hooks()) {
    auto new_grad = hook(autograd_meta->acc_grad());
    if (new_grad) { autograd_meta->set_acc_grad(new_grad); }
  }


  return Maybe<void>::Ok();
}

 

这就是autograd的基本处理过程,更多细节需要大家去查看OneFlow的autograd模块的代码。

 

六、Reference

 

本文主要参考资料是OneFlow的官方代码:

 

https://github.com/Oneflow-Inc/oneflow

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注