[源碼解析] Pytorch 如何實現后向傳播 (1)---- 調用引擎
0x00 摘要
本系列將通過大概十篇左右文章來分析 PyTorch 的自動微分功能如何實現。本文是后向傳播的第一篇，介紹調用流程：如何從 Python 代碼進入到 C++ autograd 引擎。
系列前幾篇連接如下：
[源碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)
[源碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)
[源碼解析] PyTorch如何實現前向傳播(3) --- 具體實現
0x01 前文回顧
我們首先從三個角度來看看前向傳播和后向傳播的聯系。
1.1 訓練過程
我們首先回憶一下訓練過程。
神經網絡 (NN) 是對某些輸入數據執行的嵌套函數的集合。這些函數由參數 （由權重和偏差組成）定義，這些參數在 PyTorch 中存儲在張量中。訓練 NN 分兩步進行：
-
前向傳播：在前向傳播中，神經網絡對正確的輸出做出最好的猜測。它通過它的每個函數運行輸入數據來做出這個猜測。
-
反向傳播：在反向傳播中，神經網絡根據其猜測中的誤差成比例地調整其參數。它通過從輸出向后遍歷，收集關於函數參數（梯度）的誤差導數，並使用梯度下降優化參數來實現這一點。
1.2 例子
其次，我們回憶一下前文示例。
def train_loop(model, optimizer, iterations):
for _ in range(iterations):
optimizer.zero_grad()
output = model(input) # 前向傳播
loss = criterion(output, target) # 計算損失
loss.backward() # 反向傳播
optimizer.step()
前向計算結束之后，我們已經得到了計算圖的依賴關系，於是可以開始進行后向傳播了。我們需要從 backward 開始分析。
1.3 源碼剖析
從前文我們可以看到，前向計算函數 sub_Tensor 針對前向計算結果 result 做了如下配置：
- 如何知道調用反向計算 ：result 就是前向計算的結果，result 之中有
autograd_meta_，其是一個 DifferentiableViewMeta 類型，DifferentiableViewMeta 的 grad_fn_ 就是反向計算的梯度函數。grad_fn_ 指向了 SubBackward0。 - 反向傳播如何計算 ：調用 SubBackward0 計算。
- SubBackward0 的輸入 ：得到了前向計算的輸出 result（其會在反向傳播時候作為輸入變量，就是設定到了 SubBackward0.input_metadata_ 之上）。
- SubBackward0 的輸出 ：構建了
next_edges_作為其反向傳播時候的輸出邊。根據next_edges_就能得到反向傳導圖了。
既然梳理了前向傳播與后向傳播的關系，我們接下來就看看如何進入到后向傳播環節。
0x02 Python 調用過程
2.1 調用
我們首先來到了 torch/_tensor.py，這里有兩個函數可以計算梯度，我們選取 backward 來看看。
def backward(self, gradient=None, retain_graph=None, create_graph=False, inputs=None):
r"""Computes the gradient of current tensor w.r.t. graph leaves.
"""
if has_torch_function_unary(self):
return handle_torch_function(
Tensor.backward,
(self,),
self,
gradient=gradient,
retain_graph=retain_graph,
create_graph=create_graph,
inputs=inputs)
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
然后來到了 torch/autograd/__init__.py。這里 backward 主要邏輯是：
- 利用輸入參數來構建輸入張量和梯度張量 。
- 使用 _make_grads 把 grad_tensors 中的元素重新組織成tuple(list(torch.Tensor, ...))的形式。
- 然后利用 Variable._execution_engine.run_backward 執行后向傳播。
def backward(
tensors: _TensorOrTensors,
grad_tensors: Optional[_TensorOrTensors] = None,
retain_graph: Optional[bool] = None,
create_graph: bool = False,
grad_variables: Optional[_TensorOrTensors] = None,
inputs: Optional[_TensorOrTensors] = None,
) -> None:
r"""Computes the sum of gradients of given tensors with respect to graph
leaves.
"""
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'.")
if inputs is not None and len(inputs) == 0:
raise RuntimeError("'inputs' argument to backward() cannot be empty.")
# 利用輸入參數來構建輸入張量和梯度張量
tensors = (tensors,) if isinstance(tensors, torch.Tensor) else tuple(tensors)
inputs = (inputs,) if isinstance(inputs, torch.Tensor) else \
tuple(inputs) if inputs is not None else tuple()
# _make_grads 把 grad_tensors 中的元素重新組織成tuple(list(torch.Tensor, ...))的形式
grad_tensors_ = _tensor_or_tensors_to_tuple(grad_tensors, len(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, inputs,
allow_unreachable=True, accumulate_grad=True) # allow_unreachable flag
Variable._execution_engine.run_backward 這里開始進入了C++世界。
Python + C++
|
|
|
backward |
+ |
| |
| |
| |
v |
Variable._execution_engine.run_backward +---------->
|
|
|
|
|
|
|
|
|
|
|
|
|
|
+
2.2 引擎
torch/autograd/variable.py 文件之中，生成了 _execution_engine。
from torch._C import _ImperativeEngine as ImperativeEngine
Variable._execution_engine = ImperativeEngine()
從 torch/_C/__init__.pyi.in 我們可以看到，C++世界應該去python_engine.cpp尋找答案。
# Defined in torch/csrc/autograd/python_engine.cpp
class _ImperativeEngine:
0x03 c++世界
進入 C++ 世界之后，我們放慢下腳步，先回憶一下支撐系統，否則會因為太復雜而繞暈。
3.1 支撐系統
3.1.1 Edge
Edge 通過function，input_nr 的配對來表示圖中的邊。
using tensor_list = std::vector<at::Tensor>;
using variable_list = std::vector<Variable>;
using edge_list = std::vector<Edge>;
using saved_variable_list = std::vector<SavedVariable>;
using IndexRange = std::pair<size_t, size_t>;
/// 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_) {}
/// The function this `Edge` points to.
std::shared_ptr<Node> function; // 本邊指向的Node
/// The identifier of a particular input to the function.
uint32_t input_nr; //指定本Edge在后向傳播之中是function的第幾個輸入
};
}} // namespace torch::autograd
3.1.2 Edge 相關函數
torch/csrc/autograd/function.h 這里是邊相關的函數。都是 Node 類的函數。
void set_next_edge(size_t index, Edge edge) {
update_topological_nr(edge);
next_edges_[index] = std::move(edge);
}
void add_next_edge(Edge edge) {
update_topological_nr(edge);
next_edges_.push_back(std::move(edge));
}
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);
}
}
const Edge& next_edge(size_t index) const noexcept {
return next_edges_[index];
}
const edge_list& next_edges() const noexcept {
return next_edges_;
}
edge_list& next_edges() noexcept {
return next_edges_;
}
uint32_t num_outputs() const noexcept {
return next_edges_.size();
}
torch/csrc/jit/runtime/graph_executor.cpp 之中也有一些edge相關函數。
void addOutputForTensor(const at::Tensor& tensor) {
auto v = Variable(tensor);
add_next_edge(
v.defined() ? torch::autograd::impl::gradient_edge(v)
: autograd::Edge{});
}
void addOutputForIValue(const IValue& value) {
if (value.isTensorList()) {
for (const at::Tensor tensor : value.toTensorList()) {
addOutputForTensor(tensor);
}
} else if (value.isTensor()) {
addOutputForTensor(value.toTensor());
} else {
// We could have None passed here via `Optional[Tensor]`
add_next_edge(autograd::Edge{});
}
}
gradient_edge 在前文和本文下面會用到，就是利用一個Variable的梯度和前向傳播的輸出來構建一個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.
// self.grad_fn() 這里觸發了一個調用
if (const auto& gradient = self.grad_fn()) { // 這是一個中間節點，gradient 是一個Function，比如可以得到一個SubBackward0實例
return Edge(gradient, self.output_nr()); // self.output_nr() 表示本Edge是function的第n個輸入。前向傳播時候的第 n 個輸出在反向傳播時候就是第 n 個輸入。
} else {
return Edge(grad_accumulator(self), 0); // 這是一個葉子節點，所以生成一個AccumulateGrad，0表示本Edge是function的第一個輸入
}
}
3.1.3 Python 擴展
我們接下來介紹 Python 擴展。一般來說，人們不會用C直接編寫Python模塊，而是直接寫C模塊，然后包裝一下讓Python可以直接調用，過程大致是：
- C 語言引入 Python.h 頭文件。
- 編寫封裝函數，該函數處理從 Python 世界傳入的參數。
- C 語言實現功能邏輯。
- 把 C 語言的返回值包裝成 Python 對象。
- 在 PyMethodDef 結構體中注冊所需要的函數。
- 在初始化方法中注冊模塊名。
- 把 C 源文件編譯成鏈接庫以供Python使用。
PyMethodDef 的定義如下：
typedef PyObject *(*PyCFunction)(PyObject *, PyObject *);
struct PyMethodDef {
const char *ml_name; /* The name of the built-in function/method */
PyCFunction ml_meth; /* The C function that implements it */
int ml_flags; /* Combination of METH_xxx flags, which mostly
describe the args expected by the C func */
const char *ml_doc; /* The __doc__ attribute, or NULL */
};
typedef struct PyMethodDef PyMethodDef;
3.2 引入
3.2.1 初始化
在 torch/csrc/Module.cpp 之中，initModule 會進行 C++ 世界的初始化。這是一個龐大的函數，對於本文，我們只關注 THPFunction_initModule 和 THPEngine_initModule，省略了眾多代碼。
PyObject* initModule() {
......
ASSERT_TRUE(THPFunction_initModule(module));
ASSERT_TRUE(THPEngine_initModule(module));
......
}
3.2.1.1 初始化繼承體系
初始化時候，THPFunction_initModule(module) 創建了torch._C._FunctionBase。
bool THPFunction_initModule(PyObject *module)
{
if (PyType_Ready(&THPFunctionType) < 0)
return false;
Py_INCREF(&THPFunctionType);
// 創建了`torch._C._FunctionBase`
PyModule_AddObject(module, "_FunctionBase", (PyObject *)&THPFunctionType);
return true;
}
而在torch/autograd/function.py中，有以下兩個類以torch._C._FunctionBase為基類：
class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin))
class BackwardCFunction(_C._FunctionBase, _ContextMethodMixin, _HookMixin)
這個Function繼承體系就構成了DAG的基礎。
3.2.2.2 初始化引擎
THPEngine_initModule(module) 創建了torch._C._EngineBase，_EngineBase這個類負責動態圖執行之前的預處理，_EngineBase會將torch.autograd的backward之類的請求預處理后送給真正的Engine去執行。
PyObject* initModule() {
......
ASSERT_TRUE(THPVariable_initModule(module));
ASSERT_TRUE(THPFunction_initModule(module));
ASSERT_TRUE(THPEngine_initModule(module)); // 這里初始化引擎
}
THPEngine_initModule 通過函數PyModule_AddObject 把 THPEngineType 這個對象注冊到模塊 module（一個PyObject類型） 之中，命名為 _ImperativeEngine。對應的就是 Python端的 _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);
// 為 Python 注冊了引擎
PyModule_AddObject(module, "_ImperativeEngine", (PyObject *)&THPEngineType);
set_default_engine_stub(python::PythonEngine::get_python_engine);
return true;
}
THPEngineType 定義如下，可以看出來，生成的實例是 "torch._C._EngineBase"。
PyTypeObject THPEngineType = {
PyVarObject_HEAD_INIT(nullptr, 0)
"torch._C._EngineBase", /* tp_name */
sizeof(THPEngine), /* tp_basicsize */
0, /* tp_itemsize */
nullptr, /* tp_dealloc */
0, /* tp_vectorcall_offset */
nullptr, /* tp_getattr */
nullptr, /* tp_setattr */
nullptr, /* tp_reserved */
nullptr, /* tp_repr */
nullptr, /* tp_as_number */
nullptr, /* tp_as_sequence */
nullptr, /* tp_as_mapping */
nullptr, /* tp_hash */
nullptr, /* tp_call */
nullptr, /* tp_str */
nullptr, /* tp_getattro */
nullptr, /* tp_setattro */
nullptr, /* tp_as_buffer */
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
nullptr, /* tp_doc */
nullptr, /* tp_traverse */
nullptr, /* tp_clear */
nullptr, /* tp_richcompare */
0, /* tp_weaklistoffset */
nullptr, /* tp_iter */
nullptr, /* tp_iternext */
THPEngine_methods, /* tp_methods */
nullptr, /* tp_members */
nullptr, /* tp_getset */
nullptr, /* tp_base */
nullptr, /* tp_dict */
nullptr, /* tp_descr_get */
nullptr, /* tp_descr_set */
0, /* tp_dictoffset */
nullptr, /* tp_init */
nullptr, /* tp_alloc */
THPEngine_new /* tp_new */
};
3.2.3 與Python世界聯系起來
既然 C++ 的引擎已經和 Python 的引擎聯系了起來，我們再看看引擎的具體函數。
對於torch._C._EngineBase，其成員函數是 THPEngine_methods。THPEngine_methods 的類型就是我們前面介紹的 PyMethodDef，用來進行 Python 拓展。這里定義了 run_backward，queue_callback 和 is_checkpoint_valid。我們回憶一下，run_backward 就是 Python世界的切入點。
static struct PyMethodDef THPEngine_methods[] = {
{(char*)"run_backward",
castPyCFunctionWithKeywords(THPEngine_run_backward), // 與Python對應
METH_VARARGS | METH_KEYWORDS, nullptr},
{(char*)"queue_callback", THPEngine_queue_callback, METH_O, nullptr},
{(char*)"is_checkpoint_valid", THPEngine_is_checkpoint_valid, METH_NOARGS, nullptr},
{nullptr}
};
按照前面 PyMethodDef 的定義有："run_backward" 是方法名字，THPEngine_run_backward 是對應的C語言方法。所以，Python 世界的 Variable._execution_engine.run_backward 就對應了 THPEngine_run_backward。
Python + C++
|
| initModule
| +
| |
| |
| |
| v
backward | THPEngine_initModule
+ | +
| | |
| | |
| | |
v | v
Variable._execution_engine.run_backward | PyModule_AddObject(module, "_ImperativeEngine", &THPEngineType)
+ | +
| | |
| | |
| | v
| |
| | +----------------------------------------------------------+
v | | module |
| | |
+-------------------------+ | | +---------------------------------------------------+ |
| _ImperativeEngine | | | | _ImperativeEngine | |
Variable._execution_engine +---> | | | | | | |
| | | | | +----------------------------------------------+ | |
| | | | | | THPEngine_methods | | |
| | | | | | | | |
| | | | | | | | |
| run_backward +-----------------------------> "run_backward" : THPEngine_run_backward | | |
| | | | | | | | |
| | | | | +----------------------------------------------+ | |
+-------------------------+ | | | | |
| | +---------------------------------------------------+ |
| | |
+ +----------------------------------------------------------+
手機如下：
於是我們要在 C++ 世界分析 THPEngine_run_backward。
3.3 C++引擎入口
THPEngine_run_backward 是 C++ 引擎的入口，位於：torch/csrc/autograd/python_engine.cpp。
主要邏輯如下：
-
首先，是通過函數
PyArg_ParseTupleAndKeywords對輸入的參數重新解析，並賦值給新定義的變量:- 新的變量為：
tensors，grad_tensors，keep_graph，create_graph，inputs以及allow_unreachable。比如 inputs就是一個vector。 - python世界中的輸入是 torch.autograd.backward(tensors, grad_tensors)，這些參數分別轉換被成了C++世界中的tensors和grad_tensors變量。這兩個變量在C++中的類型是PyObject，並且size為1。PyObject是任何python對象的基類，在本方法之中，tensors和grad_tensors 其實是THPVariable類的實例。
- 新的變量為：
-
從輸入獲取輸入張量和梯度張量，主要是檢查tensors和grad_tensors的變量類型以及tuple size是否一致。
-
依據輸入構建了三個變量
edge_list roots，output_edges和variable_list grads，這三個分別是反向傳播（求導）的起始點，模型最終輸出的邊信息和梯度。- roots是包含有前向傳播輸出節點的 gradient_edge()（即輸出節點的
(grad_fn_, 0)）的 vector。需要注意，grad_fn_是 Node 的派生類。 - grads 是前向傳播產生的梯度，如果沒有配置，則初始化為(tensor(1.),)。
- output_edges 是依據前向傳播輸入節點 inputs 構建的后向傳播輸出邊。
- roots是包含有前向傳播輸出節點的 gradient_edge()（即輸出節點的
-
調用outputs = engine.execute(roots, grads, keep_graph, create_graph, output_edges)，正式進入反向傳播引擎。
具體代碼如下：
// Implementation of torch._C._EngineBase.run_backward
PyObject *THPEngine_run_backward(PyObject *self, PyObject *args, PyObject *kwargs)
{
HANDLE_TH_ERRORS
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;
unsigned char accumulate_grad = 0; // Indicate whether to accumulate grad into leaf Tensors or capture
const char *accepted_kwargs[] = { // NOLINT
"tensors", "grad_tensors", "keep_graph", "create_graph", "inputs",
"allow_unreachable", "accumulate_grad", nullptr
};
// 對輸入的參數重新解析並賦值給新定義的變量tensors,grad_tensors等等，比如 inputs就是一個vector
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OObb|Obb", (char**)accepted_kwargs,
&tensors, &grad_tensors, &keep_graph, &create_graph, &inputs, &allow_unreachable, &accumulate_grad))
return nullptr;
// 從輸入獲取輸入張量和梯度張量，主要是檢查tensors和grad_tensors的變量類型以及tuple size是否一致。
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);
// The user either called autograd.backward(...) or autograd.grad(...) to get here
bool backward_api_called = accumulate_grad;
// 我們回憶一下定義
// using variable_list = std::vector<Variable>;
// using edge_list = std::vector<Edge>;
edge_list roots; // 就是反向傳播的起點（根節點）
roots.reserve(num_tensors);
variable_list grads; // 就是反向傳播的梯度
grads.reserve(num_tensors);
// 依據輸入來配置roots和grads
for (int i = 0; i < num_tensors; i++) {
// tensors是輸入節點，即前向傳播圖的輸出
PyObject *_tensor = PyTuple_GET_ITEM(tensors, i);
THPUtils_assert(THPVariable_Check(_tensor), "element %d of tensors "
// 得到 gradient_edge = Edge(grad_fn(), output_nr())
auto gradient_edge = torch::autograd::impl::gradient_edge(variable);
roots.push_back(std::move(gradient_edge)); // root增加一個Edge
PyObject *grad = PyTuple_GET_ITEM(grad_tensors, i);
if (THPVariable_Check(grad)) {
const Variable& grad_var = THPVariable_Unpack(grad);
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); // 增加一個梯度
}
}
// 構建一個輸出Edge列表
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);
const auto& tensor = THPVariable_Unpack(input);
const auto output_nr = tensor.output_nr();
auto grad_fn = tensor.grad_fn();
if (!grad_fn) {
// 獲取 grad_accumulator，用來判斷是否是葉子節點
grad_fn = torch::autograd::impl::try_get_grad_accumulator(tensor);
}
if (!grad_fn) {
// NOTE [ Autograd Unreachable Input ]
// Since input has no grad_accumulator, its guaranteed to be unreachable.
// We initialize an edge pointing to a non-nullptr Node so nodes in the graph
// (e.g., mul when an operand is scalar) that have edges pointing to nullptr
// don't get erroneously assigned `needed = True` in exec_info.
// 說明是葉子節點
output_edges.emplace_back(std::make_shared<Identity>(), 0);
} else {
// 是中間節點
output_edges.emplace_back(grad_fn, output_nr);
}
}
}
// 現在，roots是包含有(前向傳播輸出節點的grad_fn_, 0)的vector。
// grads 是前向傳播產生的梯度，如果沒有配置，則初始化為(tensor(1.),)
// output_edges 是依據前向傳播輸入節點 input 構建的后向傳播輸出邊
variable_list outputs;
{
pybind11::gil_scoped_release no_gil;
auto& engine = python::PythonEngine::get_python_engine();
// 進入引擎執行
outputs = engine.execute(roots, grads, keep_graph, create_graph, accumulate_grad, output_edges);
}
if (!backward_api_called && 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++) {
PyTuple_SET_ITEM(py_outputs.get(), i, THPVariable_Wrap(outputs[i]));
}
return py_outputs.release();
} else {
Py_RETURN_NONE;
}
END_HANDLE_TH_ERRORS
}
我們接下來分析 THPEngine_run_backward 用到的幾個輔助函數。
3.3.1 try_get_grad_accumulator
上面代碼之中，有 grad_fn = torch::autograd::impl::try_get_grad_accumulator(tensor) 來獲取計算梯度的方法。其實是用它來判斷是否是葉子節點，只有非葉子節點 grad_accumulator_才不為空。
try_get_grad_accumulator 返回的是個指向Node對象的指針 : std::weak_ptr<Node> grad_accumulator_。就是如何計算梯度。
具體邏輯是：
- 先通過函數
get_autograd_meta返回一個AutogradMeta結構體。 - 然后訪問結構體中的成員變量
grad_accumulator_，而grad_accumulator_是一個指向類型為Node對象的std::weak_ptr指針。 - 最后通過
lock()函數創建一個std::shared_ptr來管理對象。
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;
}
}
3.3.2 gradient_edge
上面代碼之中，gradient_edge 被用來在輸入 tensor 基礎之上構建一個 Edge。
auto gradient_edge = torch::autograd::impl::gradient_edge(variable);
roots.push_back(std::move(gradient_edge)); // root增加一個Edge
gradient_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);
}
}
3.3.3 output_edges
上面代碼之中， std::vector
在拿到 grad_accumulator_ 之后，會賦予為 grad_fn，這樣就用來判斷是否為葉子節點。然后分別構建葉子節點和中間節點，放到 output_edges 之中。
if (!grad_fn) {
// NOTE [ Autograd Unreachable Input ]
// Since input has no grad_accumulator, its guaranteed to be unreachable.
// We initialize an edge pointing to a non-nullptr Node so nodes in the graph
// (e.g., mul when an operand is scalar) that have edges pointing to nullptr
// don't get erroneously assigned `needed = True` in exec_info.
output_edges.emplace_back(std::make_shared<Identity>(), 0); // 葉子節點
} else {
output_edges.emplace_back(grad_fn, output_nr); // 非葉子節點
}
我們看看構建output_edges的變量 grad_fn 和 output_nr，看看它們的由來。
grad_fn是通過 try_get_grad_accumulator 方法得到的一個指向Node對象的std::shared_ptr指針，就是如何計算梯度的操作。
output_pr 由如下設置，其最終得到的是結構體AutogradMeta中的成員變量uint32_t output_nr_。
const auto output_nr = tensor.output_nr();
emplace_back()函數向容器中中加入臨時對象， 臨時對象原地構造，沒有賦值或移動的操作。
回憶一下 Edge 的定義。所以可以看出來，emplace_back()就是使用了這些輸入生成了一個 Edge。
/// 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_) {}
/// The function this `Edge` points to.
std::shared_ptr<Node> function; // 指向的Node
/// The identifier of a particular input to the function.
uint32_t input_nr; //指定本Edge在后向傳播之中是function的第幾個輸入
};
輸入轉換如下圖，可以看出來輸入從 Python 如何進行轉換最終傳入C++引擎，以如下變量為例：
- Python 的 tensors 被轉換為 C++ 的 root。
- Python 的 grad_tensors 被轉換為 C++ 的 grads。
- Python 的 inputs 被轉換為 C++ 的 output_edges。
- 最終把這三個變量傳遞給引擎：PythonEngine.execute(roots, grads, keep_graph, create_graph, accumulate_grad, output_edges)。
backward(tensors, grad_tensors, inputs)
+ + +
| | |
Python | | |
| | |
+------------------------------------------------------------------------------------------+
| | |
C++ THPEngine_run_backward | | |
| | +-----------------------------+
| | |
| | |
| +-----------------------------+ |
v | |
| |
+------root = [(tensor_1.grad_fn_, 0),...,(tensor_n.grad_fn_, 0)] | |
| | |
| | |
| | |
| +--grads = [grad_tensor_1,...,grad_tensor_n ] <----------------------+ |
| | |
| | |
| | v
| | output_edges = [(input_1.grad_fn_, output_nr_1),...,(input_n.grad_fn_, output_nr_n)]
| | +
| +-------------------------+ |
| | |
| | |
+----------------------+ | |
| | |
v v v
PythonEngine.execute(roots, grads, keep_graph, create_graph, accumulate_grad, output_edges)
3.4 PythonEngine
前面 THPEngine_run_backward 代碼有如下，我們可以看到，THPEngine_run_backward 最終調用到了 PythonEngine 的處理邏輯。
auto& engine = python::PythonEngine::get_python_engine();
// 進入引擎執行
outputs = engine.execute(roots, grads, keep_graph, create_graph, accumulate_grad, output_edges);
3.4.1 獲取引擎
get_python_engine這里定義了一個靜態變量。整個PyTorch程序全局只維護一個Engine實例，也就是PythonEngine實例。
Engine& PythonEngine::get_python_engine() {
static PythonEngine engine;
// This is "probably" thread-safe because the flag is set in a fork handler
// before any threads are created, and this function is only called with the
// GIL held. However, using fork + threads is playing with fire so this is
// more of a "best effort" thing. For example, if the fork occurs while the
// backwards threads hold a lock, we'll probably deadlock in the engine
// destructor.
if (_reinitialize_engine) {
engine.release_workers();
engine.~PythonEngine();
new (&engine) torch::autograd::python::PythonEngine();
_reinitialize_engine = false;
}
return engine;
}
3.4.2 定義
所以我們來看看PythonEngine 定義。PythonEngine 是 Engine 的派生類，相當於封裝了一下。主要是針對python世界的特點做了一些定制，比如：PythonEngine子類重寫了父類的execute，把C++異常翻譯為Python異常的功能，核心工作還是由Engine基類來完成：
struct PythonEngine : public Engine {
static Engine& get_python_engine();
~PythonEngine() override;
void thread_init(int device,
const std::shared_ptr<ReadyQueue>& ready_queue,
bool should_increment) 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,
bool accumulate_grad,
const edge_list& outputs = {}) override;
std::shared_ptr<at::ivalue::Future> execute_with_graph_task(
const std::shared_ptr<GraphTask>& graph_task,
std::shared_ptr<Node> graph_root,
InputBuffer&& input_buffer) override;
std::unique_ptr<AnomalyMetadata> make_anomaly_metadata() override;
private:
PythonEngine();
};
execute 代碼如下，於是從下文開始，我們要看看 Engine 是如何運作的。
variable_list PythonEngine::execute(
const edge_list& roots,
const variable_list& inputs,
bool keep_graph,
bool create_graph,
bool accumulate_grad,
const edge_list& outputs) {
try {
return Engine::execute(roots, inputs, keep_graph, create_graph, accumulate_grad, outputs);
} catch (python_error& e) {
e.restore();
throw;
}
}
目前邏輯拓展如下：
backward(tensors, grad_tensors, inputs)
+ + +
| | |
Python | | |
| | |
+------------------------------------------------------------------------------------------+
| | |
C++ THPEngine_run_backward | | |
| | +-----------------------------+
| | |
| | |
| +-----------------------------+ |
v | |
| |
+------root = [(tensor_1.grad_fn_, 0),...,(tensor_n.grad_fn_, 0)] | |
| | |
| | |
| | |
| +--grads = [grad_tensor_1,...,grad_tensor_n ] <----------------------+ |
| | |
| | |
| | v
| | output_edges = [(input_1.grad_fn_, output_nr_1),...,(input_n.grad_fn_, output_nr_n)]
| | +
| +-------------------------+ |
| | |
| | |
+----------------------+ | |
| | |
v v v
PythonEngine.execute(roots, grads, keep_graph, create_graph, accumulate_grad, output_edges)
+ + + +
| | | |
| | | |
v v v v
Engine::execute(roots, inputs, keep_graph, create_graph, accumulate_grad, outputs)
手機如下：
3.5 另一調用途徑
最后，我們再插入一個 run_backward 進行分析。
run_backward 位於 torch/csrc/autograd/autograd.cpp。這里應該是專門為了 C++ 世界直接調用的需要，與我們之前通過 Python 迂回調用不同。
void backward(
const variable_list& tensors,
const variable_list& grad_tensors,
c10::optional<bool> retain_graph,
bool create_graph,
const variable_list& inputs) {
variable_list gradients = _make_grads(tensors, grad_tensors);
if (!retain_graph) {
retain_graph = create_graph;
}
run_backward(tensors, gradients, retain_graph.value(), create_graph, inputs, /*allow_unused=*/true, /*accumulate_grad=*/true);
}
variable_list grad(
const variable_list& outputs,
const variable_list& inputs,
const variable_list& grad_outputs,
c10::optional<bool> retain_graph,
bool create_graph,
bool allow_unused) {
variable_list gradients = _make_grads(outputs, grad_outputs);
if (!retain_graph) {
retain_graph = create_graph;
}
return run_backward(
outputs, gradients, retain_graph.value(), create_graph, inputs, allow_unused, /*accumulate_grad=*/false);
}
run_backward 最后也調用了 Engine::get_default_engine().execute。
variable_list run_backward(
const variable_list& outputs,
const variable_list& grad_outputs,
bool keep_graph,
bool create_graph,
const variable_list& inputs,
bool allow_unused,
bool accumulate_grad) {
size_t num_tensors = outputs.size();
edge_list roots;
roots.reserve(num_tensors);
for (size_t i = 0; i < num_tensors; i++) {
const Variable& output = outputs[i];
auto gradient_edge = impl::gradient_edge(output);
roots.push_back(std::move(gradient_edge));
}
edge_list output_edges;
if (!inputs.empty()) {
size_t num_inputs = inputs.size();
output_edges.reserve(num_inputs);
for (size_t i = 0; i < num_inputs; ++i) {
const Variable& input = inputs[i];
const auto output_nr = input.output_nr();
auto grad_fn = input.grad_fn();
if (!grad_fn) {
grad_fn = impl::try_get_grad_accumulator(input);
}
if (!grad_fn) {
// See NOTE [ Autograd Unreachable Input ] for details
output_edges.emplace_back(std::make_shared<Identity>(), 0);
} else {
output_edges.emplace_back(grad_fn, output_nr);
}
}
}
// 調用了引擎代碼
variable_list grad_inputs = Engine::get_default_engine().execute(
roots, grad_outputs, keep_graph, create_graph, accumulate_grad, output_edges);
// check if grad_inputs contains None or not base on the allow_unused flag
if (!inputs.empty() && !allow_unused) {
size_t num_inputs = inputs.size();
for (size_t i = 0; i < num_inputs; ++i) {
TORCH_CHECK(
grad_inputs[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.");
}
}
return grad_inputs;
}
至此，調用過程分析完畢，其核心就是調用引擎函數進行處理，所以下一篇我們來開始分析引擎。