【PyTorch】聊聊 backward 背后的代码

Example:
>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> optimizer.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optimizer.step() 

def backward(self, gradient=None, retain_graph=None, create_graph=False):
    r"""Computes the gradient of current tensor w.r.t. graph leaves.

    The graph is differentiated using the chain rule. If the tensor is
    non-scalar (i.e. its data has more than one element) and requires
    gradient, the function additionally requires specifying ``gradient``.
    It should be a tensor of matching type and location, that contains
    the gradient of the differentiated function w.r.t. ``self``.

    This function accumulates gradients in the leaves - you might need to
    zero them before calling it.

    Arguments:
        gradient (Tensor or None): Gradient w.r.t. the
            tensor. If it is a tensor, it will be automatically converted
            to a Tensor that does not require grad unless ``create_graph`` is True.
            None values can be specified for scalar Tensors or ones that
            don't require grad. If a None value would be acceptable then
            this argument is optional.
        retain_graph (bool, optional): If ``False``, the graph used to compute
            the grads will be freed. Note that in nearly all cases setting
            this option to True is not needed and often can be worked around
            in a much more efficient way. Defaults to the value of
            ``create_graph``.
        create_graph (bool, optional): If ``True``, graph of the derivative will
            be constructed, allowing to compute higher order derivative
            products. Defaults to ``False``.
    """
    torch.autograd.backward(self, gradient, retain_graph, create_graph) 

# ...

from .variable import Variable

# ... 

def _make_grads(outputs, grads):
    new_grads = []
    for out, grad in zip(outputs, grads):
        if isinstance(grad, torch.Tensor):
            if not out.shape == grad.shape:
                # raise RuntimeError ...
            new_grads.append(grad)
        elif grad is None:
            if out.requires_grad:
                if out.numel() != 1:
                    # raise RuntimeError ...
            else:
                new_grads.append(None)
        else:
            # raise TypeError ...
    return tuple(new_grads)


def backward(tensors, grad_tensors=None, retain_graph=None, create_graph=False, grad_variables=None):
    r"""Computes the sum of gradients of given tensors w.r.t. graph leaves.

    The graph is differentiated using the chain rule. If any of ``tensors``
    are non-scalar (i.e. their data has more than one element) and require
    gradient, then the Jacobian-vector product would be computed, in this
    case the function additionally requires specifying ``grad_tensors``.
    It should be a sequence of matching length, that contains the "vector"
    in the Jacobian-vector product, usually the gradient of the differentiated
    function w.r.t. corresponding tensors (``None`` is an acceptable value for
    all tensors that don't need gradient tensors).

    This function accumulates gradients in the leaves - you might need to zero
    them before calling it.
    """
    if grad_variables is not None:
        warnings.warn("'grad_variables' is deprecated. Use 'grad_tensors' instead.")
        if grad_tensors is None:
            grad_tensors = grad_variables
        else:
            raise RuntimeError("'grad_tensors' and 'grad_variables' (deprecated) "
                               "arguments both passed to backward(). Please only "
                               "use 'grad_tensors'.")

    tensors = (tensors,) if isinstance(tensors, torch.Tensor) else tuple(tensors)

    if grad_tensors is None:
        grad_tensors = [None] * len(tensors)
    elif isinstance(grad_tensors, torch.Tensor):
        grad_tensors = [grad_tensors]
    else:
        grad_tensors = list(grad_tensors)

    grad_tensors = _make_grads(tensors, grad_tensors)
    if retain_graph is None:
        retain_graph = create_graph

    Variable._execution_engine.run_backward(
        tensors, grad_tensors, retain_graph, create_graph,
        allow_unreachable=True)  # allow_unreachable flag

# ...

if not torch._C._autograd_init():
    raise RuntimeError("autograd initialization failed")

import torch
from torch._six import with_metaclass


class VariableMeta(type):
    def __instancecheck__(cls, other):
        return isinstance(other, torch.Tensor)


class Variable(with_metaclass(VariableMeta, torch._C._LegacyVariableBase)):
    pass


from torch._C import _ImperativeEngine as ImperativeEngine
Variable._execution_engine = ImperativeEngine()

bool THPEngine_initModule(PyObject *module)
{
#ifndef _WIN32
  if (pthread_atfork(nullptr, nullptr, child_atfork) != 0) {
    throw std::runtime_error("unable to set pthread_atfork handler");
  }
#endif
  if (PyType_Ready(&THPEngineType) < 0)
    return false;
  Py_INCREF(&THPEngineType);
  PyModule_AddObject(module, "_ImperativeEngine", (PyObject *)&THPEngineType);
  set_default_engine_stub(get_python_engine);
  return true;
}

// ...
 
static struct PyMethodDef THPEngine_methods[] = {
  {(char*)"run_backward", (PyCFunction)(void(*)(void))THPEngine_run_backward, METH_VARARGS | METH_KEYWORDS, nullptr},
  {(char*)"queue_callback", (PyCFunction)THPEngine_queue_callback, METH_O, nullptr},
  {(char*)"is_checkpoint_valid", (PyCFunction)THPEngine_is_checkpoint_valid, METH_NOARGS, nullptr},
  {nullptr}
};

// ...

PyTypeObject THPEngineType = {
    // ...
    "torch._C._EngineBase",                      /* tp_name */
    // ...
    THPEngine_methods,                           /* tp_methods */
    // ...
};

// ...

// Implementation of torch._C._EngineBase.run_backward
PyObject *THPEngine_run_backward(THPEngine *self, PyObject *args, PyObject *kwargs)
{
    HANDLE_TH_ERRORS
    _maybe_reinitialize_engine_after_fork();
    PyObject *tensors = nullptr;
    PyObject *grad_tensors = nullptr;
    unsigned char keep_graph = 0;
    unsigned char create_graph = 0;
    PyObject *inputs = nullptr;
    unsigned char allow_unreachable = 0;
    const char *accepted_kwargs[] = {
            "tensors", "grad_tensors", "keep_graph", "create_graph", "inputs",
            "allow_unreachable", nullptr
    };
    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OObb|Ob", (char**)accepted_kwargs,
                &tensors, &grad_tensors, &keep_graph, &create_graph, &inputs, &allow_unreachable))
        return nullptr;

    // ...

    END_HANDLE_TH_ERRORS
}

// Implementation of torch._C._EngineBase.run_backward
PyObject *THPEngine_run_backward(THPEngine *self, PyObject *args, PyObject *kwargs)
{
    // ...

    THPUtils_assert(PyTuple_Check(tensors), "tensors argument is expected to "
            "be a tuple, but got %s", THPUtils_typename(tensors));
    THPUtils_assert(PyTuple_Check(grad_tensors), "grad_tensors argument is "
            "expected to be a tuple, but got %s", THPUtils_typename(grad_tensors));

    Py_ssize_t num_tensors = PyTuple_GET_SIZE(tensors);
    Py_ssize_t num_gradients = PyTuple_GET_SIZE(grad_tensors);
    THPUtils_assert(num_tensors == num_gradients, "got %ld tensors and %ld "
            "gradients", num_tensors, num_gradients);

    // ...
}

// Implementation of torch._C._EngineBase.run_backward
PyObject *THPEngine_run_backward(THPEngine *self, PyObject *args, PyObject *kwargs)
{
    // ...

    edge_list roots;
    roots.reserve(num_tensors);
    variable_list grads;
    grads.reserve(num_tensors);
    for (int i = 0; i < num_tensors; i++) {
        PyObject *_tensor = PyTuple_GET_ITEM(tensors, i);
        THPUtils_assert(THPVariable_Check(_tensor), "element %d of tensors "
                "tuple is not a Tensor", i);
        auto& variable = ((THPVariable*)_tensor)->cdata;
        auto gradient_edge = torch::autograd::impl::gradient_edge(variable);
        THPUtils_assert(gradient_edge.function,
                "element %d of tensors does not require grad and does not have a grad_fn", i);
        roots.push_back(std::move(gradient_edge));

        PyObject *grad = PyTuple_GET_ITEM(grad_tensors, i);
        if (THPVariable_Check(grad)) {
            const Variable& grad_var = ((THPVariable*)grad)->cdata;
            if (grad_var.has_names()) {
                TORCH_WARN(
                        "Autograd was passed a named grad tensor with dims ", grad_var.names(),
                        ". Autograd does not yet support named tensor semantics, so all names ",
                        "will be ignored. In practice all computed gradients will still be correct "
                        "according to regular tensor semantics.");
            }
            grads.push_back(grad_var);
        } else {
            THPUtils_assert(grad == Py_None,
                    "element %d of gradients tuple is not a Tensor or None", i);
            THPUtils_assert(!variable.requires_grad(),
                    "element %d of gradients tuple is None, but the corresponding Tensor requires grad");
        }
    }

    // ...
}

// Implementation of torch._C._EngineBase.run_backward
PyObject *THPEngine_run_backward(THPEngine *self, PyObject *args, PyObject *kwargs)
{
    // ...

    std::vector<Edge> output_edges;
    if (inputs != nullptr) {
        int num_inputs = PyTuple_GET_SIZE(inputs);
        output_edges.reserve(num_inputs);
        for (int i = 0; i < num_inputs; ++i) {
            PyObject *input = PyTuple_GET_ITEM(inputs, i);
            THPUtils_assert(THPVariable_Check(input),
                    "all inputs have to be Tensors, but got %s", THPUtils_typename(input));
            THPVariable *input_var = (THPVariable*)input;
            const auto output_nr = input_var->cdata.output_nr();
            auto grad_fn = input_var->cdata.grad_fn();
            if (!grad_fn) {
                    grad_fn = torch::autograd::impl::try_get_grad_accumulator(input_var->cdata);
            }
            THPUtils_assert(input_var->cdata.requires_grad(),
                    "One of the differentiated Tensors does not require grad");
            if (!grad_fn) {
                output_edges.emplace_back();
            } else {
                output_edges.emplace_back(grad_fn, output_nr);
            }
        }
    }

    // ...
}

// python_variable.h

// Python object that backs torch.autograd.Variable
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
struct THPVariable {
    PyObject_HEAD
    // Payload
    torch::autograd::Variable cdata;
    // Hooks to be run on backwards pass (corresponds to Python attr
    // '_backwards_hooks', set by 'register_hook')
    PyObject* backward_hooks = nullptr;
};

// =======================================================

// VariableTypeManual.cpp

int64_t output_nr(const Tensor & self) {
    if (impl::get_autograd_meta(self)) {
        return impl::get_autograd_meta(self)->output_nr_;
    } else {
        return 0;
    }
}

// =======================================================

// variable.h

struct TORCH_API AutogradMeta : public c10::AutogradMetaInterface {
    // ...
    // The "output number" of this variable; e.g., if this variable
    // was the second output of a function, then output_nr == 1.
    // We use this to make sure we can setup the backwards trace
    // correctly when this variable is passed to another function.
    uint32_t output_nr_;
    // ...
};

// =======================================================

// edge.h

/// Represents a particular input of a function.
struct Edge {
    Edge() noexcept : function(nullptr), input_nr(0) {}
    Edge(std::shared_ptr<Node> function_, uint32_t input_nr_) noexcept
        : function(std::move(function_)), input_nr(input_nr_) {}
    // ...
};

// Implementation of torch._C._EngineBase.run_backward
PyObject *THPEngine_run_backward(THPEngine *self, PyObject *args, PyObject *kwargs)
{
    // ...

    variable_list outputs;
    {
        pybind11::gil_scoped_release no_gil;
        outputs = engine.execute(roots, grads, keep_graph, create_graph, output_edges);
    }

    if (inputs != nullptr) {
        int num_inputs = PyTuple_GET_SIZE(inputs);
        THPObjectPtr py_outputs {PyTuple_New(num_inputs)};
        if (!py_outputs) return nullptr;
        for (int i = 0; i < num_inputs; i++) {
            THPUtils_assert(allow_unreachable || outputs[i].defined(), "One of the "
                                            "differentiated Tensors appears to not have been used "
                                            "in the graph. Set allow_unused=True if this is the "
                                            "desired behavior.");
            PyTuple_SET_ITEM(py_outputs.get(), i, THPVariable_Wrap(outputs[i]));
        }
        return py_outputs.release();
    } else {
        Py_RETURN_NONE;
    }

    // ...
}

// variable.h
// ...
namespace torch { namespace autograd {
struct Node; 
struct AutogradMeta;
struct DifferentiableViewMeta;
using Variable = at::Tensor;
namespace impl {
    // ...
    TORCH_API AutogradMeta* get_autograd_meta(const Variable&);
    // ...
    TORCH_API std::shared_ptr<Node> try_get_grad_accumulator(const Variable&);
    // ...
}
// ...
struct TORCH_API AutogradMeta : public c10::AutogradMetaInterface {
    // ...
    std::weak_ptr<Node> grad_accumulator_; 
    // ...
};

// =================================

// variable.cpp
// ...
namespace torch {
namespace autograd {
// ...
namespace impl {
    // ...
    std::shared_ptr<Node> try_get_grad_accumulator(const Variable& self) {
        if (get_autograd_meta(self)) {
            return get_autograd_meta(self)->grad_accumulator_.lock();
        } else {
            return nullptr;
        }
    }
    // ...
    AutogradMeta* get_autograd_meta(const Variable& self) {
        // NB: could return null
        TORCH_CHECK(self.defined(), "cannot call get_autograd_meta() on undefined tensor");
        return static_cast<AutogradMeta*>(self.unsafeGetTensorImpl()->autograd_meta());
    }
    // ...
}
// ...

static torch::autograd::python::PythonEngine engine; 

// python_engine.h

namespace torch { namespace autograd { namespace python {

struct PythonEngine : public Engine {
  void thread_init(int device) override;
  void thread_on_exception(
      std::shared_ptr<GraphTask>& graph_task,
      const std::shared_ptr<Node>& fn,
      std::exception& e) override;
  variable_list execute(
      const edge_list& roots,
      const variable_list& inputs,
      bool keep_graph,
      bool create_graph,
      const edge_list& outputs = {}) override;

  variable_list execute_with_graph_task(
      std::shared_ptr<GraphTask> graph_task,
      std::shared_ptr<Node> graph_root) override;
  std::unique_ptr<AnomalyMetadata> make_anomaly_metadata() override;
};

}}} // namespace torch::autograd::python

// ================================================

// python_engine.cpp

variable_list PythonEngine::execute(
    const edge_list& roots,
    const variable_list& inputs,
    bool keep_graph,
    bool create_graph,
    const edge_list& outputs) {
  try {
    return Engine::execute(roots, inputs, keep_graph, create_graph, outputs);
  } catch (python_error& e) {
    e.restore();
    throw;
  }
}

variable_list Engine::execute(const edge_list& roots, 
                              const variable_list& inputs, 
                              bool keep_graph, 
                              bool create_graph, 
                              const edge_list& outputs)
// 调用 ↓
variable_list Engine::execute_with_graph_task(std::shared_ptr<GraphTask> graph_task,
                                              std::shared_ptr<Node> graph_root)
// 调用 ↓
void Engine::thread_main(const std::shared_ptr<GraphTask>& graph_task,
                         bool reentrant_thread)
// 调用 ↓
void Engine::evaluate_function(std::shared_ptr<GraphTask>& graph_task,
                               Node* func, 
                               InputBuffer& inputs)
// 调用 ↓ （这个函数不是 Engine 结构体中的方法）
variable_list call_function(std::shared_ptr<GraphTask>& graph_task,
                            Node* func,
                            InputBuffer& inputBuffer) {
    // ...
    auto& fn = *func;
    auto inputs = call_pre_hooks(fn, InputBuffer::variables(std::move(inputBuffer)));
    variable_list outputs = fn(std::move(inputs));
    // ...
    if(has_post_hooks){
        // NOLINTNEXTLINE(bugprone-use-after-move)
        return call_post_hooks(fn, std::move(outputs), inputs);
    }
    return outputs;
}

// =================================

static variable_list call_pre_hooks(Node& fn, 
                                    variable_list inputs) {
    for (const auto& hook : fn.pre_hooks()) {
        inputs = (*hook)(inputs);
    }
    return inputs;
}

static variable_list call_post_hooks(Node& fn, 
                                     variable_list outputs, 
                                     const variable_list& inputs) {
    for (const auto& hook : fn.post_hooks()) {
        outputs = (*hook)(outputs, inputs);
    }
    return outputs;
}

#pragma once

#include <vector>
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <ATen/Tensor.h>

// A hook that's called on gradients

namespace torch { namespace autograd {

using Variable = at::Tensor;
using variable_list = std::vector<Variable>;

struct TORCH_API FunctionPreHook {
    virtual ~FunctionPreHook();
    virtual variable_list operator()(const variable_list& grads) = 0;
};

struct TORCH_API FunctionPostHook {
    virtual ~FunctionPostHook();
    virtual variable_list operator()(
        const variable_list& outputs /* grad_inputs */,
        const variable_list& inputs /* grad_outputs */) = 0;
};

}} // namespace torch::autograd