在上一篇文章中，我们讨论了自动微分的理论基础，并回顾了 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 包含用于计算给定函数的雅可比向量积、海森矩阵和其他梯度相关计算的组件。其余的文件包含额外的组件，如梯度检查器、异常检测和 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++中定义，将 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 与输入张量链接起来。执行 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 对象，这是一个特殊的 Node 对象，当在正向传播过程中执行 apply 时，它会调用 Python 解释器执行提供的 Python 函数。请注意，在代码中， cdata 是实际属于图中的 Node 对象。 ctx 是传递给 Python forward / backward 函数的对象，它被用户函数和 PyTorch 都用来存储与自动微分相关的信息。
正如常规的 C++ 函数一样，我们也会在 collect_next_edges 中调用 grad_fn 来跟踪输入的 unpack_input 对象，但这是在 @3# 中完成的：
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 设置到新创建的变量中。为此，它调用另一个定义在不同文件中的 _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 负责在输出张量中设置 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 引擎如何执行这些图。
