在上一篇文章中，我们回顾了自动微分的理论基础，并讨论了 PyTorch 中的实现。在本篇文章中，我们将展示 PyTorch 中涉及图创建和执行的部分。为了理解以下内容，请阅读 @ezyang 关于 PyTorch 内部原理的精彩博客文章。

Autograd 组件

首先，让我们看看 autograd 的不同组件所在位置

tools/autograd：在这里我们可以找到我们上一篇文章中看到的导数定义derivatives.yaml，几个 Python 脚本和一个名为templates的文件夹。这些脚本和模板在构建时用于生成 yaml 文件中指定的导数的 C++ 代码。此外，这里的脚本还为常规 ATen 函数生成封装器，以便可以构建计算图。

torch/autograd：此文件夹是 autograd 组件所在的位置，可以直接从 python 中使用。在function.py中，我们找到了实际的 `torch.autograd.Function` 定义，这是一个供用户根据文档在 python 中编写自己的可微分函数所使用的类。functional.py包含用于函数式计算给定函数的雅可比向量积、Hessian 和其他与梯度相关的计算的组件。其余文件包含其他组件，例如梯度检查器、异常检测和 autograd 分析器。

torch/csrc/autograd：这是图创建和执行相关代码的所在地。所有这些代码都用 C++ 编写，因为它是一个需要极高性能的关键部分。在这里，我们有几个文件实现了引擎、元数据存储以及所有所需的组件。此外，我们还有几个以 `python_` 开头的文件，它们的主要职责是允许在 autograd 引擎中使用 python 对象。

图创建

之前，我们描述了计算图的创建。现在，我们将看到 PyTorch 如何通过引用实际代码库来创建这些图。

图 1：增强计算图示例

这一切都始于我们的 Python 代码中，当我们请求一个张量需要梯度时。

>>> x = torch.tensor([0.5, 0.75], requires_grad=True)

当在张量创建时设置 `required_grad` 标志时，c10 将分配一个 `AutogradMeta` 对象，该对象用于保存图信息。

void TensorImpl::set_requires_grad(bool requires_grad) {
  ...
  if (!autograd_meta_)
    autograd_meta_ = impl::GetAutogradMetaFactory()->make();
    autograd_meta_->set_requires_grad(requires_grad, this);
}

`AutogradMeta` 对象在torch/csrc/autograd/variable.h中定义如下

struct TORCH_API AutogradMeta : public c10::AutogradMetaInterface {
  std::string name_;

  Variable grad_;
  std::shared_ptr<Node> grad_fn_;
  std::weak_ptr<Node> grad_accumulator_;
  // other fields and methods
  ...
};

此结构中最重要的字段是 `grad_` 中的计算梯度和指向 `grad_fn` 函数的指针，该函数将由引擎调用以生成实际梯度。此外，还有一个梯度累加器对象，用于将此张量涉及的所有不同梯度相加，我们将在图执行中看到。

图、节点和边。

现在，当我们调用一个以该张量为参数的可微函数时，关联的元数据将被填充。假设我们调用一个在 ATen 中实现的常规 torch 函数。例如我们之前博客文章中的乘法示例。结果张量有一个名为 `grad_fn` 的字段，它本质上是指向将用于计算该操作梯度的函数的指针。

>>> x = torch.tensor([0.5, 0.75], requires_grad=True)
>>> v = x[0] * x[1]
>>> v
tensor(0.3750, grad_fn=<MulBackward0>)

在这里我们看到张量的 `grad_fn` 具有 `MulBackward0` 值。此函数与在derivatives.yaml文件中写入的函数相同，其 C++ 代码由 `tools/autograd` 中的所有脚本自动生成。其自动生成的源代码可在 `torch/csrc/autograd/generated/Functions.cpp` 中找到。

variable_list MulBackward0::apply(variable_list&& grads) {
  std::lock_guard<std::mutex> lock(mutex_);

  IndexRangeGenerator gen;
  auto self_ix = gen.range(1);
  auto other_ix = gen.range(1);
  variable_list grad_inputs(gen.size());
  auto& grad = grads[0];
  auto self = self_.unpack();
  auto other = other_.unpack();
  bool any_grad_defined = any_variable_defined(grads);
  if (should_compute_output({ other_ix })) {
    auto grad_result = any_grad_defined ? (mul_tensor_backward(grad, self, other_scalar_type)) : Tensor();
    copy_range(grad_inputs, other_ix, grad_result);
  }
  if (should_compute_output({ self_ix })) {
    auto grad_result = any_grad_defined ? (mul_tensor_backward(grad, other, self_scalar_type)) : Tensor();
    copy_range(grad_inputs, self_ix, grad_result);
  }
  return grad_inputs;
}

`grad_fn` 对象继承自 `TraceableFunction` 类，它是 `Node` 的子孙，只设置了一个属性以启用用于调试和优化目的的跟踪。图根据定义具有节点和边，因此这些函数确实是计算图的节点，它们通过使用 `Edge` 对象连接在一起，以便稍后进行图遍历。

`Node` 的定义可以在 torch/csrc/autograd/function.h 文件中找到。

struct TORCH_API Node : std::enable_shared_from_this<Node> {
 ...
 /// Evaluates the function on the given inputs and returns the result of the
  /// function call.
  variable_list operator()(variable_list&& inputs) {
  ...
  }

protected:
  /// Performs the `Node`'s actual operation.
  virtual variable_list apply(variable_list&& inputs) = 0;
  …
  edge_list next_edges_;

我们基本上可以看到它重写了 `operator ()` 来执行实际函数调用，以及一个纯虚函数 `apply`。正如我们在上面的 `MulBackward0` 示例中看到的那样，自动生成的函数会重写这个 `apply` 方法。最后，节点还包含一个边列表，以实现图连接。

Edge 对象用于连接 `Node`，其实现非常简单。

struct Edge {
  ...
  /// The function this `Edge` points to.
  std::shared_ptr<Node> function;
  /// The identifier of a particular input to the function.
  uint32_t input_nr;
};

它只需要一个函数指针（连接边的实际 `grad_fn` 对象）和一个作为边 ID 的输入编号。

链接节点

当我们调用两个张量的乘法运算时，我们就进入了自动生成代码的领域。我们在 `tools/autograd` 中看到的所有脚本都填充了一系列模板，这些模板封装了 ATen 中的可微函数。这些函数包含在正向传播期间构建反向图的代码。

gen_variable_type.py 脚本负责编写所有这些包装代码。该脚本在 PyTorch 构建过程中从 tools/autograd/gen_autograd.py 调用，并将自动生成的函数包装器输出到 `torch/csrc/autograd/generated/`。

让我们来看看张量乘法生成的函数是什么样子。代码已简化，但在从源代码编译 PyTorch 时，可以在 `torch/csrc/autograd/generated/VariableType_4.cpp` 文件中找到它。

at::Tensor mul_Tensor(c10::DispatchKeySet ks, const at::Tensor & self, const at::Tensor & other) {
  ...
  auto _any_requires_grad = compute_requires_grad( self, other );
  std::shared_ptr<MulBackward0> grad_fn;
  if (_any_requires_grad) {
    // Creates the link to the actual grad_fn and links the graph for backward traversal
    grad_fn = std::shared_ptr<MulBackward0>(new MulBackward0(), deleteNode);
    grad_fn->set_next_edges(collect_next_edges( self, other ));
    ...
  }
  …
  // Does the actual function call to ATen
  auto _tmp = ([&]() {
    at::AutoDispatchBelowADInplaceOrView guard;
    return at::redispatch::mul(ks & c10::after_autograd_keyset, self_, other_);
  })();

  auto result = std::move(_tmp);
    if (grad_fn) {
       // Connects the result to the graph
      set_history(flatten_tensor_args( result ), grad_fn);
  }
  ...
  return result;
}

让我们逐行了解此代码中最重要的部分。首先，`grad_fn` 对象通过以下代码创建：`grad_fn = std::shared_ptr(new MulBackward0(), deleteNode);`。

在 `grad_fn` 对象创建后，用于连接节点的边通过 `grad_fn->set_next_edges(collect_next_edges( self, other ));` 调用创建。

struct MakeNextFunctionList : IterArgs<MakeNextFunctionList> {
  edge_list next_edges;
  using IterArgs<MakeNextFunctionList>::operator();
  void operator()(const Variable& variable) {
    if (variable.defined()) {
      next_edges.push_back(impl::gradient_edge(variable));
    } else {
      next_edges.emplace_back();
    }
  }
  void operator()(const c10::optional<Variable>& variable) {
    if (variable.has_value() && variable->defined()) {
      next_edges.push_back(impl::gradient_edge(*variable));
    } else {
      next_edges.emplace_back();
    }
  }
};

template <typename... Variables>
edge_list collect_next_edges(Variables&&... variables) {
  detail::MakeNextFunctionList make;
  make.apply(std::forward<Variables>(variables)...);
  return std::move(make.next_edges);
}

给定一个输入变量（它只是一个常规张量），`collect_next_edges` 将通过调用`impl::gradient_edge` 来创建 `Edge` 对象。

 Edge gradient_edge(const Variable& self) {
    // If grad_fn is null (as is the case for a leaf node), we instead
    // interpret the gradient function to be a gradient accumulator, which will
    // accumulate its inputs into the grad property of the variable. These
    // nodes get suppressed in some situations, see "suppress gradient
    // accumulation" below. Note that only variables which have `requires_grad =
    // True` can have gradient accumulators.
    if (const auto& gradient = self.grad_fn()) {
      return Edge(gradient, self.output_nr());
    } else {
      return Edge(grad_accumulator(self), 0);
    }
  }

要理解边缘如何工作，假设一个早期执行的函数产生了两个输出张量，它们都设置了 `grad_fn`，每个张量还有一个 `output_nr` 属性，表示它们返回的顺序。在为当前 `grad_fn` 创建边缘时，将为每个输入变量创建一个 `Edge` 对象。边缘将指向变量的 `grad_fn`，并将跟踪 `output_nr` 以建立在遍历图时使用的 ID。如果输入变量是“叶子”，即它们不是由任何可微函数产生的，则它们没有设置 `grad_fn` 属性。默认情况下会设置一个特殊的函数，称为梯度累加器，如上面的代码片段所示。

边缘创建后，当前正在创建的 `grad_fn` 图节点对象将使用 `set_next_edges` 函数保存它们。这就是连接 `grad_fn` 的方式，从而生成计算图。

 void set_next_edges(edge_list&& next_edges) {
    next_edges_ = std::move(next_edges);
    for(const auto& next_edge : next_edges_) {
      update_topological_nr(next_edge);
    }
  }

现在，函数的前向传播将执行，执行后 `set_history` 将输出张量连接到 `grad_fn` 节点。

inline void set_history(
    at::Tensor& variable,
    const std::shared_ptr<Node>& grad_fn) {
  AT_ASSERT(grad_fn);
  if (variable.defined()) {
    // If the codegen triggers this, you most likely want to add your newly added function
    // to the DONT_REQUIRE_DERIVATIVE list in tools/autograd/gen_variable_type.py
    TORCH_INTERNAL_ASSERT(isDifferentiableType(variable.scalar_type()));
    auto output_nr =
        grad_fn->add_input_metadata(variable);
    impl::set_gradient_edge(variable, {grad_fn, output_nr});
  } else {
    grad_fn->add_input_metadata(Node::undefined_input());
  }
}

`set_history` 调用`set_gradient_edge`，它只是将 `grad_fn` 和 `output_nr` 复制到张量拥有的 `AutogradMeta` 对象中。

 void set_gradient_edge(const Variable& self, Edge edge) {
    auto* meta = materialize_autograd_meta(self);
    meta->grad_fn_ = std::move(edge.function);
    meta->output_nr_ = edge.input_nr;
    // For views, make sure this new grad_fn_ is not overwritten unless it is necessary
    // in the VariableHooks::grad_fn below.
    // This logic is only relevant for custom autograd Functions for which multiple
    // operations can happen on a given Tensor before its gradient edge is set when
    // exiting the custom Function.
    auto diff_view_meta = get_view_autograd_meta(self);
    if (diff_view_meta && diff_view_meta->has_bw_view()) {
      diff_view_meta->set_attr_version(self._version());
    }
  }

这个张量现在将作为另一个函数的输入，并且上述步骤将全部重复。查看下面的动画以了解图是如何创建的。

图 2：显示图创建的动画

在图中注册 Python 函数

我们已经看到了 autograd 如何为 ATen 中包含的函数创建图。但是，当我们在 Python 中定义可微分函数时，它们也包含在图中！

一个 autograd python 定义的函数看起来像下面这样

class Exp(torch.autograd.Function):
     @staticmethod
     def forward(ctx, i):
         result = i.exp()
         ctx.save_for_backward(result)
         return result

     @staticmethod
     def backward(ctx, grad_output):
         result, = ctx.saved_tensors
         return grad_output * result

# Call the function
Exp.apply(torch.tensor(0.5, requires_grad=True))
# Outputs: tensor(1.6487, grad_fn=<ExpBackward>)

在上面的代码片段中，autograd 在创建图时检测到了我们的 python 函数。所有这些都归功于 `Function` 类。让我们看看当我们调用 `apply` 时会发生什么。

`apply` 在 `torch._C._FunctionBase` 类中定义，但此类不存在于 python 源代码中。`_FunctionBase` 是通过使用 python C API 将 C 函数连接到一个 python 类中来在 C++ 中定义的。我们正在寻找一个名为 `THPFunction_apply` 的函数。

PyObject *THPFunction_apply(PyObject *cls, PyObject *inputs)
{
  
  // Generates the graph node
  THPObjectPtr backward_cls(PyObject_GetAttrString(cls, "_backward_cls"));
  if (!backward_cls) return nullptr;
  THPObjectPtr ctx_obj(PyObject_CallFunctionObjArgs(backward_cls, nullptr));
  if (!ctx_obj) return nullptr;
  THPFunction* ctx = (THPFunction*)ctx_obj.get();

  auto cdata = std::shared_ptr<PyNode>(new PyNode(std::move(ctx_obj)), deleteNode);
  ctx->cdata = cdata;

  // Prepare inputs and allocate context (grad fn)
  // Unpack inputs will collect the edges
  auto info_pair = unpack_input<false>(inputs);
  UnpackedInput& unpacked_input = info_pair.first;
  InputFlags& input_info = info_pair.second;

   // Initialize backward function (and ctx)
  bool is_executable = input_info.is_executable;
  cdata->set_next_edges(std::move(input_info.next_edges));
  ctx->needs_input_grad = input_info.needs_input_grad.release();
  ctx->is_variable_input = std::move(input_info.is_variable_input);

  // Prepend ctx to input_tuple, in preparation for static method call
  auto num_args = PyTuple_GET_SIZE(inputs);
  THPObjectPtr ctx_input_tuple(PyTuple_New(num_args + 1));
  if (!ctx_input_tuple) return nullptr;
  Py_INCREF(ctx);
  PyTuple_SET_ITEM(ctx_input_tuple.get(), 0, (PyObject*)ctx);
  for (int i = 0; i < num_args; ++i) {
    PyObject *arg = PyTuple_GET_ITEM(unpacked_input.input_tuple.get(), i);
    Py_INCREF(arg);
    PyTuple_SET_ITEM(ctx_input_tuple.get(), i + 1, arg);
  }

  // Call forward
  THPObjectPtr tensor_outputs;
  {
    AutoGradMode grad_mode(false);
    THPObjectPtr forward_fn(PyObject_GetAttrString(cls, "forward"));
    if (!forward_fn) return nullptr;
    tensor_outputs = PyObject_CallObject(forward_fn, ctx_input_tuple);
    if (!tensor_outputs) return nullptr;
  }

  // Here is where the outputs gets the tensors tracked
  return process_outputs(cls, cdata, ctx, unpacked_input, inputs, std::move(tensor_outputs),
                         is_executable, node);
  END_HANDLE_TH_ERRORS
}

尽管由于所有 Python API 调用，这段代码一开始难以阅读，但它本质上与我们为 ATen 看到的自动生成的前向函数执行相同的事情

创建 `grad_fn` 对象。收集边缘以将当前的 `grad_fn` 与输入张量连接起来。执行函数 `forward`。将创建的 `grad_fn` 分配给输出张量元数据。

`grad_fn` 对象创建于

  // Generates the graph node
  THPObjectPtr backward_cls(PyObject_GetAttrString(cls, "_backward_cls"));
  if (!backward_cls) return nullptr;
  THPObjectPtr ctx_obj(PyObject_CallFunctionObjArgs(backward_cls, nullptr));
  if (!ctx_obj) return nullptr;
  THPFunction* ctx = (THPFunction*)ctx_obj.get();

  auto cdata = std::shared_ptr<PyNode>(new PyNode(std::move(ctx_obj)), deleteNode);
  ctx->cdata = cdata;

基本上，它要求 Python API 获取指向可以执行用户编写函数的 Python 对象的指针。然后它将其封装到 `PyNode` 对象中，`PyNode` 是一个特殊的 `Node` 对象，它在正向传播期间执行 `apply` 时使用提供的 Python 函数调用 Python 解释器。请注意，在代码中，`cdata` 是图中实际的 `Node` 对象。`ctx` 是传递给 Python `forward`/`backward` 函数的对象，它用于存储与自动梯度相关的信息，供用户函数和 PyTorch 使用。

与常规 C++ 函数一样，我们也调用 `collect_next_edges` 来跟踪输入的 `grad_fn` 对象，但这在 `unpack_input` 中完成。

template<bool enforce_variables>
std::pair<UnpackedInput, InputFlags> unpack_input(PyObject *args) {
  ...
  flags.next_edges = (flags.is_executable ? collect_next_edges(unpacked.input_vars) : edge_list());
  return std::make_pair(std::move(unpacked), std::move(flags));
}

此后，通过 `cdata->set_next_edges(std::move(input_info.next_edges));` 将边分配给 `grad_fn`，并通过 Python 解释器 C API 调用前向函数。

一旦输出张量从前向传播返回，它们将在 `process_outputs` 函数内部进行处理并转换为变量。

PyObject* process_outputs(PyObject *op_obj, const std::shared_ptr<PyNode>& cdata,
                          THPFunction* grad_fn, const UnpackedInput& unpacked,
                          PyObject *inputs, THPObjectPtr&& raw_output, bool is_executable,
                          torch::jit::Node* node) {
  ...
  _wrap_outputs(cdata, grad_fn, unpacked.input_vars, raw_output, outputs, is_executable);
  _trace_post_record(node, op_obj, unpacked.input_vars, outputs, is_inplace, unpack_output);
  if (is_executable) {
    _save_variables(cdata, grad_fn);
  } ...
  return outputs.release();
}

在这里，`_wrap_outputs` 负责将前向输出的 `grad_fn` 设置为新创建的 `grad_fn`。为此，它调用了另一个在不同文件中定义的 `_wrap_outputs` 函数，因此这里的过程有点令人困惑。

static void _wrap_outputs(const std::shared_ptr<PyNode>& cdata, THPFunction *self,
    const variable_list &input_vars, PyObject *raw_output, PyObject *outputs, bool is_executable)
{
  auto cdata_if_executable = is_executable ? cdata : nullptr;
 ...

  // Wrap only the tensor outputs.
  // This calls csrc/autograd/custom_function.cpp
  auto wrapped_outputs = _wrap_outputs(input_vars, non_differentiable, dirty_inputs, raw_output_vars, cdata_if_executable);
...
}

被调用的 `_wrap_outputs` 负责设置输出张量中的自动梯度元数据

std::vector<c10::optional<Variable>> _wrap_outputs(const variable_list &input_vars,
  const std::unordered_set<at::TensorImpl*> &non_differentiable,
  const std::unordered_set<at::TensorImpl*> &dirty_inputs,
  const at::ArrayRef<c10::optional<Variable>> raw_outputs,
  const std::shared_ptr<Node> &cdata) {


  std::unordered_set<at::TensorImpl*> inputs;
  …
  // Sets the grad_fn and output_nr of an output Variable.
  auto set_history = [&](Variable& var, uint32_t output_nr, bool is_input, bool is_modified,
                         bool is_differentiable) {
    // Lots of checks
    if (!is_differentiable) {
     ...
    } else if (is_input) {
      // An input has been returned, but it wasn't modified. Return it as a view
      // so that we can attach a new grad_fn to the Variable.
      // Run in no_grad mode to mimic the behavior of the forward.
      {
        AutoGradMode grad_mode(false);
        var = var.view_as(var);
      }
      impl::set_gradient_edge(var, {cdata, output_nr});
    } else if (cdata) {
      impl::set_gradient_edge(var, {cdata, output_nr});
    }
  };

这就是 `set_gradient_edge` 被调用的地方，也是用户编写的 Python 函数如何与其关联的反向函数一起包含在计算图中的方式！

总结

这篇博文旨在概述 PyTorch 如何构建我们在上一篇博文中讨论的实际计算图。下一篇文章将讨论自动梯度引擎如何执行这些图。