在上一篇博文中，我们回顾了自动微分的理论基础，并审视了 PyTorch 中的实现。在本篇博文中，我们将展示 PyTorch 中涉及图创建和执行的部分。为了理解以下内容，请阅读 @ezyang 关于 PyTorch 内部机制的精彩博文。
Autograd 组件
首先，让我们看看 autograd 的不同组件位于何处
tools/autograd：在这里，我们可以找到我们在上一篇博文 derivatives.yaml 中看到的导数定义，以及几个 Python 脚本和一个名为 templates 的文件夹。这些脚本和模板在构建时用于根据 yaml 文件中指定的导数生成 C++ 代码。此外，这里的脚本还为常规的 ATen 函数生成包装器，以便可以构建计算图。
torch/autograd：该文件夹存放着可以直接从 Python 中使用的 autograd 组件。在 function.py 中，我们可以找到 torch.autograd.Function 的实际定义，用户可以使用该类根据文档编写自己的可微分函数（在 Python 中）。functional.py 包含用于函数式计算给定函数的雅可比向量积、Hessian 矩阵和其他梯度相关计算的组件。其余文件包含其他组件，例如梯度检查器、异常检测和 autograd 分析器。
torch/csrc/autograd：这是图创建和执行相关代码的所在地。所有这些代码都用 C++ 编写，因为这是需要极高性能的关键部分。这里有几个文件实现了引擎、元数据存储以及所有必需的组件。此外，我们还有一些文件名以 python_ 开头的文件，它们的主要职责是允许在 autograd 引擎中使用 Python 对象。
图创建
此前，我们描述了计算图的创建过程。现在，我们将参考实际代码库，看看 PyTorch 如何创建这些图。
图 1：增强计算图示例
这一切始于我们的 Python 代码，我们在其中要求一个 tensor 计算梯度。
>>> x = torch.tensor([0.5, 0.75], requires_grad=True)
在创建 tensor 时设置 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);
}
根据 torch/csrc/autograd/variable.h 中的定义，AutogradMeta 对象如下所示
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 函数的指针，引擎将调用该函数来生成实际梯度。此外，还有一个梯度累加器对象，用于累加此 tensor 所涉及的所有不同梯度，这将在图执行部分看到。
图、节点和边。
现在，当我们调用一个以该 tensor 作为参数的可微分函数时，关联的元数据将被填充。假设我们调用一个在 ATen 中实现的常规 torch 函数。例如，就像我们在上一篇博文示例中看到的乘法运算。结果 tensor 有一个名为 grad_fn 的字段，它本质上是指向用于计算该操作梯度的函数的指针。
>>> x = torch.tensor([0.5, 0.75], requires_grad=True)
>>> v = x[0] * x[1]
>>> v
tensor(0.3750, grad_fn=<MulBackward0>)
这里我们看到 tensor 的 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。
连接节点
当我们调用两个 tensor 的乘法运算时，我们就进入了自动生成代码的领域。我们在 tools/autograd 中看到的所有脚本都填充了一系列模板，用于包装 ATen 中的可微分函数。这些函数包含在正向传播期间构建反向图的代码。
gen_variable_type.py 脚本负责编写所有这些包装代码。该脚本在 PyTorch 构建过程中由 tools/autograd/gen_autograd.py 调用，它将把自动生成的函数包装器输出到 torch/csrc/autograd/generated/。
让我们看看 tensor 乘法生成的函数是什么样子。代码已简化，但在从源代码编译 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
创建 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);
}
给定一个输入变量（它只是一个常规 tensor），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);
}
}
为了理解边的工作原理，假设一个早期执行的函数生成了两个输出 tensor，它们的 grad_fn 都已设置，每个 tensor 还有一个 output_nr 属性，表示它们的返回顺序。在为当前 grad_fn 创建边时，每个输入变量将创建一个 Edge 对象。边将指向变量的 grad_fn，并跟踪 output_nr 以建立图遍历时使用的 ID。如果输入变量是“叶子”变量（即它们不是由任何可微分函数生成的），则它们没有设置 grad_fn 属性。默认情况下会设置一个名为梯度累加器的特殊函数，如上面的代码片段所示。
创建边后，当前正在创建的 grad_fn 图 Node 对象将使用 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 将输出 tensor 连接到 grad_fn Node。
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 复制到 tensor 拥有的 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());
}
}
现在，此 tensor 将作为另一个函数的输入，上述所有步骤都将重复进行。查看下面的动画，了解图是如何创建的。
图 2：展示图创建过程的动画
在图中注册 Python 函数
我们已经了解了 autograd 如何为 ATen 中包含的函数创建图。但是，当我们在 Python 中定义自己的可微分函数时，它们也会被包含在图中！
一个用 Python 定义的 autograd 函数如下所示
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++ 中定义的，用于将 C 函数连接到单个 Python 类中。我们正在寻找一个名为 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 与输入 tensor 连接起来。执行 forward 函数。将创建的 grad_fn 分配给输出 tensor 的元数据。
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 对象中，该对象是一个专门的 Node 对象，在正向传播期间执行 apply 时，它会使用提供的 Python 函数调用 Python 解释器。请注意，在代码中 cdata 是图中的实际 Node 对象。ctx 是传递给 Python forward/backward 函数的对象，它由用户函数和 PyTorch 都用于存储 autograd 相关信息。
与常规 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 调用 forward 函数。
从正向传播返回输出 tensor 后，它们会在 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 负责将 forward 输出的 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 是负责在输出 tensor 中设置 autograd 元数据的函数。
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 如何构建我们在上一篇博文中讨论过的实际计算图。下一篇文章将探讨 autograd 引擎如何执行这些图。
