乐趣区

关于人工智能:Autograd解析｜OneFlow学习笔记

jiezi

撰文|月踏

更新|赵露阳

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

1

一个求梯度的小例子

先看上面这个简略的例子：

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 模块。

2

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 模式，这是因为 graph 动态图模式下，计算图是提前编译好的，无需手动通过.backward() 调用。flow.autograd.backward() 会调用 oneflow/api/python/autograd/autograd.cpp 中导出的 backward 办法：

ONEFLOW_API_PYBIND11_MODULE("autograd", m) {m.def("backward", &Backward);
  m.def("grad", &Grad);
}

从 pybind 定义来看，这外面总共导出了两个接口（autograd.backward 和 autograd.grad）。其中，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;}
  std::shared_ptr<one::TensorTuple> gradients = JUST(CheckAndInitOutGrads(outputs, out_grads));
  JUST(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，返回一个 autograd 引擎（AutogradEngine）对象的指针：

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

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


图 2

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

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);
  JUST(graph_task.ComputeDependencies());
  JUST(graph_task.Apply(/*save_grad_for_leaf=*/true));
  return Maybe<void>::Ok();}

这就真正进入了 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）

GraphFunctionNode::GraphFunctionNode(const std::string& name,
                                     const std::shared_ptr<BackwardFunction>& backward_fn,
                                     const TensorTuple& inputs, const TensorTuple& outputs)
    : FunctionNode(name, backward_fn) {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:L86），外面会创立这个反向函数：

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));


  std::shared_ptr<OpExprGradClosure> grad_closure(nullptr);
  if (requires_grad && !LazyMode::is_enabled()) {grad_closure = JUST(op_expr.GetOrCreateOpGradClosure());
    auto backward_fn = std::make_shared<BackwardFunction>();
    backward_fn->body = [=](const TensorTuple& out_grads, TensorTuple* in_grads,
                            bool create_graph) -> Maybe<void> {autograd::AutoGradMode mode(create_graph);
      JUST(grad_closure->Apply(out_grads, in_grads));
      return Maybe<void>::Ok();};
    backward_fn->status = [=]() { return grad_closure->state()->SavedTensors().size() > 0; };
    JUST(GetThreadLocalAutogradEngine()->AddNode(op_expr.op_type_name() + "_backward", backward_fn,
                                                 inputs, outputs));
  }
  ...
  return Maybe<void>::Ok();}

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

Maybe<FunctionNode> GraphAutogradEngine::AddNode(
    const std::string& name, const std::shared_ptr<BackwardFunction>& backward_fn,
    const TensorTuple& inputs, TensorTuple* outputs) {
  // Firstly push function_node of tensor in stack which is leaf and requires_grad
  for (const std::shared_ptr<Tensor>& in_tensor : inputs) {if (in_tensor->is_leaf() && in_tensor->requires_grad()) {if (!in_tensor->grad_fn_node()) {JUST(AddAccumulateFunctionNode(in_tensor)); }
    }
  }


  std::shared_ptr<FunctionNode> func_node =
      std::make_shared<GraphFunctionNode>(name, backward_fn, inputs, *outputs);
  for (const std::shared_ptr<Tensor>& 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 接口来获取。

4

反向图的执行流程

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

...
JUST(graph_task.ComputeDependencies());
JUST(graph_task.Apply(/*save_grad_for_leaf=*/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:L287）

Maybe<void> GraphTask::Apply(bool save_grad_for_leaf) {
  std::queue<FunctionNode*> queue;
  for (FunctionNode* 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 (/*bool not_ready_to_apply=*/!(JUST(node->Apply(create_graph_)))) {continue;}
    if (save_grad_for_leaf) {JUST(node->AccGrad4LeafTensor(create_graph_)); }
    JUST(node->AccGrad4RetainGradTensor());
    node->ReleaseOutTensorArgs();
    if (!retain_graph_) {node->ReleaseData(); }


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

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

  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_中：

Maybe<bool> FunctionNode::Apply(bool create_graph) {
  ...
  JUST(backward_fn_->body(output_grads, &input_grads, create_graph));
  for (int i = 0; i < input_meta_data_.size(); ++i) {if (input_grads.at(i)) {
       ...
       JUST(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) {autograd::AutoGradMode mode(autograd_mode);
  auto current_grad = JUST(autograd_meta->current_grad()->GetAccTensor({}));
  if (!current_grad) {return Maybe<void>::Ok(); }
  if (autograd_meta->acc_grad()) {
    ...
    DevVmDepObjectConsumeModeGuard guard(DevVmDepObjectConsumeMode::NONE);
    const auto& output = JUST(functional::Add(autograd_meta->acc_grad(), current_grad, /*alpha=*/1,
                                              /*inplace=*/autograd_meta->is_grad_acc_inplace()));
    JUST(autograd_meta->set_acc_grad(output));
  } else {JUST(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) {JUST(autograd_meta->set_acc_grad(new_grad)); }
  }


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

（* 特别感谢共事 yinggang 两头的各种答疑解惑。本文次要参考代码：
https://github.com/Oneflow-In…*）

欢送下载体验 OneFlow v0.7.0 最新版本：
https://github.com/Oneflow-In…

退出移动版