在之前的文章中，我们回顾了自动微分的理论基础，并回顾了 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 的纯虚函数。自动生成的函数覆盖了此 apply 方法，正如我们在上面的 MulBackward0 示例中看到的那样。最后，节点还具有边列表，以启用图连接。
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 脚本负责编写所有这些包装代码。此脚本从 tools/autograd/gen_autograd.py 在 pytorch 构建过程中调用，它将自动生成的函数包装器输出到 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
创建 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 在 C++ 中使用 python C API 定义，以将 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 的 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 调用正向函数。
一旦输出张量从正向传递返回，它们将在 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 负责在输出张量中设置 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 引擎如何执行这些图。
