[源碼解析] Pytorch 如何實(shí)現(xiàn)后向傳播 (1)---- 調(diào)用引擎
[源碼解析] Pytorch 如何實(shí)現(xiàn)后向傳播 (1)---- 調(diào)用引擎
目錄
0x00 摘要
本系列將通過(guò)大概十篇左右文章來(lái)分析 PyTorch 的自動(dòng)微分功能如何實(shí)現(xiàn)。本文是后向傳播的第一篇，介紹調(diào)用流程：如何從 Python 代碼進(jìn)入到 C++ autograd 引擎。
系列前幾篇連接如下：
[源碼解析]深度學(xué)習(xí)利器之自動(dòng)微分(3) --- 示例解讀
[源碼解析]PyTorch如何實(shí)現(xiàn)前向傳播(1) --- 基礎(chǔ)類(上)
[源碼解析]PyTorch如何實(shí)現(xiàn)前向傳播(2) --- 基礎(chǔ)類(下)
[源碼解析] PyTorch如何實(shí)現(xiàn)前向傳播(3) --- 具體實(shí)現(xiàn)
0x01 前文回顧
我們首先從三個(gè)角度來(lái)看看前向傳播和后向傳播的聯(lián)系。
1.1 訓(xùn)練過(guò)程
我們首先回憶一下訓(xùn)練過(guò)程。
神經(jīng)網(wǎng)絡(luò) (NN) 是對(duì)某些輸入數(shù)據(jù)執(zhí)行的嵌套函數(shù)的集合。這些函數(shù)由參數(shù) （由權(quán)重和偏差組成）定義，這些參數(shù)在 PyTorch 中存儲(chǔ)在張量中。訓(xùn)練 NN 分兩步進(jìn)行：
-
前向傳播：在前向傳播中，神經(jīng)網(wǎng)絡(luò)對(duì)正確的輸出做出最好的猜測(cè)。它通過(guò)它的每個(gè)函數(shù)運(yùn)行輸入數(shù)據(jù)來(lái)做出這個(gè)猜測(cè)。
-
反向傳播：在反向傳播中，神經(jīng)網(wǎng)絡(luò)根據(jù)其猜測(cè)中的誤差成比例地調(diào)整其參數(shù)。它通過(guò)從輸出向后遍歷，收集關(guān)于函數(shù)參數(shù)（梯度）的誤差導(dǎo)數(shù)，并使用梯度下降優(yōu)化參數(shù)來(lái)實(shí)現(xiàn)這一點(diǎn)。
1.2 例子
其次，我們回憶一下前文示例。
def train_loop(model, optimizer, iterations):
for _ in range(iterations):
optimizer.zero_grad()
output = model(input) # 前向傳播
loss = criterion(output, target) # 計(jì)算損失
loss.backward() # 反向傳播
optimizer.step()
前向計(jì)算結(jié)束之后，我們已經(jīng)得到了計(jì)算圖的依賴關(guān)系，于是可以開(kāi)始進(jìn)行后向傳播了。我們需要從 backward 開(kāi)始分析。
1.3 源碼剖析
從前文我們可以看到，前向計(jì)算函數(shù) sub_Tensor 針對(duì)前向計(jì)算結(jié)果 result 做了如下配置：
- 如何知道調(diào)用反向計(jì)算 ：result 就是前向計(jì)算的結(jié)果，result 之中有
autograd_meta_，其是一個(gè) DifferentiableViewMeta 類型，DifferentiableViewMeta 的 grad_fn_ 就是反向計(jì)算的梯度函數(shù)。grad_fn_ 指向了 SubBackward0。 - 反向傳播如何計(jì)算 ：調(diào)用 SubBackward0 計(jì)算。
- SubBackward0 的輸入 ：得到了前向計(jì)算的輸出 result（其會(huì)在反向傳播時(shí)候作為輸入變量，就是設(shè)定到了 SubBackward0.input_metadata_ 之上）。
- SubBackward0 的輸出 ：構(gòu)建了
next_edges_作為其反向傳播時(shí)候的輸出邊。根據(jù)next_edges_就能得到反向傳導(dǎo)圖了。
既然梳理了前向傳播與后向傳播的關(guān)系，我們接下來(lái)就看看如何進(jìn)入到后向傳播環(huán)節(jié)。
0x02 Python 調(diào)用過(guò)程
2.1 調(diào)用
我們首先來(lái)到了 torch/_tensor.py，這里有兩個(gè)函數(shù)可以計(jì)算梯度，我們選取 backward 來(lái)看看。
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)
然后來(lái)到了 torch/autograd/__init__.py。這里 backward 主要邏輯是：
- 利用輸入?yún)?shù)來(lái)構(gòu)建輸入張量和梯度張量 。
- 使用 _make_grads 把 grad_tensors 中的元素重新組織成tuple(list(torch.Tensor, ...))的形式。
- 然后利用 Variable._execution_engine.run_backward 執(zhí)行后向傳播。
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.")
# 利用輸入?yún)?shù)來(lái)構(gòu)建輸入張量和梯度張量
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
# 執(zhí)行后向傳播
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 這里開(kāi)始進(jìn)入了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++世界應(yīng)該去python_engine.cpp尋找答案。
# Defined in torch/csrc/autograd/python_engine.cpp
class _ImperativeEngine:
0x03 c++世界
進(jìn)入 C++ 世界之后，我們放慢下腳步，先回憶一下支撐系統(tǒng)，否則會(huì)因?yàn)樘珡?fù)雜而繞暈。
3.1 支撐系統(tǒng)
3.1.1 Edge
Edge 通過(guò)function，input_nr 的配對(duì)來(lái)表示圖中的邊。
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的第幾個(gè)輸入
};
}} // namespace torch::autograd
3.1.2 Edge 相關(guān)函數(shù)
torch/csrc/autograd/function.h 這里是邊相關(guān)的函數(shù)。都是 Node 類的函數(shù)。
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相關(guān)函數(shù)。
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 在前文和本文下面會(huì)用到，就是利用一個(gè)Variable的梯度和前向傳播的輸出來(lái)構(gòu)建一個(gè)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() 這里觸發(fā)了一個(gè)調(diào)用
if (const auto& gradient = self.grad_fn()) { // 這是一個(gè)中間節(jié)點(diǎn)，gradient 是一個(gè)Function，比如可以得到一個(gè)SubBackward0實(shí)例
return Edge(gradient, self.output_nr()); // self.output_nr() 表示本Edge是function的第n個(gè)輸入。前向傳播時(shí)候的第 n 個(gè)輸出在反向傳播時(shí)候就是第 n 個(gè)輸入。
} else {
return Edge(grad_accumulator(self), 0); // 這是一個(gè)葉子節(jié)點(diǎn)，所以生成一個(gè)AccumulateGrad，0表示本Edge是function的第一個(gè)輸入
}
}
3.1.3 Python 擴(kuò)展
我們接下來(lái)介紹 Python 擴(kuò)展。一般來(lái)說(shuō)，人們不會(huì)用C直接編寫(xiě)Python模塊，而是直接寫(xiě)C模塊，然后包裝一下讓Python可以直接調(diào)用，過(guò)程大致是：
- C 語(yǔ)言引入 Python.h 頭文件。
- 編寫(xiě)封裝函數(shù)，該函數(shù)處理從 Python 世界傳入的參數(shù)。
- C 語(yǔ)言實(shí)現(xiàn)功能邏輯。
- 把 C 語(yǔ)言的返回值包裝成 Python 對(duì)象。
- 在 PyMethodDef 結(jié)構(gòu)體中注冊(cè)所需要的函數(shù)。
- 在初始化方法中注冊(cè)模塊名。
- 把 C 源文件編譯成鏈接庫(kù)以供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 會(huì)進(jìn)行 C++ 世界的初始化。這是一個(gè)龐大的函數(shù)，對(duì)于本文，我們只關(guān)注 THPFunction_initModule 和 THPEngine_initModule，省略了眾多代碼。
PyObject* initModule() {
......
ASSERT_TRUE(THPFunction_initModule(module));
ASSERT_TRUE(THPEngine_initModule(module));
......
}
3.2.1.1 初始化繼承體系
初始化時(shí)候，THPFunction_initModule(module) 創(chuàng)建了torch._C._FunctionBase。
bool THPFunction_initModule(PyObject *module)
{
if (PyType_Ready(&THPFunctionType) < 0)
return false;
Py_INCREF(&THPFunctionType);
// 創(chuàng)建了`torch._C._FunctionBase`
PyModule_AddObject(module, "_FunctionBase", (PyObject *)&THPFunctionType);
return true;
}
而在torch/autograd/function.py中，有以下兩個(gè)類以torch._C._FunctionBase為基類：
class Function(with_metaclass(FunctionMeta, _C._FunctionBase, _ContextMethodMixin, _HookMixin))
class BackwardCFunction(_C._FunctionBase, _ContextMethodMixin, _HookMixin)
這個(gè)Function繼承體系就構(gòu)成了DAG的基礎(chǔ)。
3.2.2.2 初始化引擎
THPEngine_initModule(module) 創(chuàng)建了torch._C._EngineBase，_EngineBase這個(gè)類負(fù)責(zé)動(dòng)態(tài)圖執(zhí)行之前的預(yù)處理，_EngineBase會(huì)將torch.autograd的backward之類的請(qǐng)求預(yù)處理后送給真正的Engine去執(zhí)行。
PyObject* initModule() {
......
ASSERT_TRUE(THPVariable_initModule(module));
ASSERT_TRUE(THPFunction_initModule(module));
ASSERT_TRUE(THPEngine_initModule(module)); // 這里初始化引擎
}
THPEngine_initModule 通過(guò)函數(shù)PyModule_AddObject 把 THPEngineType 這個(gè)對(duì)象注冊(cè)到模塊 module（一個(gè)PyObject類型） 之中，命名為 _ImperativeEngine。對(duì)應(yīng)的就是 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 注冊(cè)了引擎
PyModule_AddObject(module, "_ImperativeEngine", (PyObject *)&THPEngineType);
set_default_engine_stub(python::PythonEngine::get_python_engine);
return true;
}
THPEngineType 定義如下，可以看出來(lái)，生成的實(shí)例是 "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世界聯(lián)系起來(lái)
既然 C++ 的引擎已經(jīng)和 Python 的引擎聯(lián)系了起來(lái)，我們?cè)倏纯匆娴木唧w函數(shù)。
對(duì)于torch._C._EngineBase，其成員函數(shù)是 THPEngine_methods。THPEngine_methods 的類型就是我們前面介紹的 PyMethodDef，用來(lái)進(jìn)行 Python 拓展。這里定義了 run_backward，queue_callback 和 is_checkpoint_valid。我們回憶一下，run_backward 就是 Python世界的切入點(diǎn)。
static struct PyMethodDef THPEngine_methods[] = {
{(char*)"run_backward",
castPyCFunctionWithKeywords(THPEngine_run_backward), // 與Python對(duì)應(yīng)
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 是對(duì)應(yīng)的C語(yǔ)言方法。所以，Python 世界的 Variable._execution_engine.run_backward 就對(duì)應(yīng)了 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 | | |
| | | | | | | | |
| | | | | +----------------------------------------------+ | |
+-------------------------+ | | | | |
| | +---------------------------------------------------+ |
| | |
+ +----------------------------------------------------------+
手機(jī)如下：
于是我們要在 C++ 世界分析 THPEngine_run_backward。
3.3 C++引擎入口
THPEngine_run_backward 是 C++ 引擎的入口，位于：torch/csrc/autograd/python_engine.cpp。
主要邏輯如下：
-
首先，是通過(guò)函數(shù)
PyArg_ParseTupleAndKeywords對(duì)輸入的參數(shù)重新解析，并賦值給新定義的變量:- 新的變量為：
tensors，grad_tensors，keep_graph，create_graph，inputs以及allow_unreachable。比如 inputs就是一個(gè)vector。 - python世界中的輸入是 torch.autograd.backward(tensors, grad_tensors)，這些參數(shù)分別轉(zhuǎn)換被成了C++世界中的tensors和grad_tensors變量。這兩個(gè)變量在C++中的類型是PyObject，并且size為1。PyObject是任何python對(duì)象的基類，在本方法之中，tensors和grad_tensors 其實(shí)是THPVariable類的實(shí)例。
- 新的變量為：
-
從輸入獲取輸入張量和梯度張量，主要是檢查tensors和grad_tensors的變量類型以及tuple size是否一致。
-
依據(jù)輸入構(gòu)建了三個(gè)變量
edge_list roots，output_edges和variable_list grads，這三個(gè)分別是反向傳播（求導(dǎo)）的起始點(diǎn)，模型最終輸出的邊信息和梯度。- roots是包含有前向傳播輸出節(jié)點(diǎn)的 gradient_edge()（即輸出節(jié)點(diǎn)的
(grad_fn_, 0)）的 vector。需要注意，grad_fn_是 Node 的派生類。 - grads 是前向傳播產(chǎn)生的梯度，如果沒(méi)有配置，則初始化為(tensor(1.),)。
- output_edges 是依據(jù)前向傳播輸入節(jié)點(diǎn) inputs 構(gòu)建的后向傳播輸出邊。
- roots是包含有前向傳播輸出節(jié)點(diǎn)的 gradient_edge()（即輸出節(jié)點(diǎn)的
-
調(diào)用outputs = engine.execute(roots, grads, keep_graph, create_graph, output_edges)，正式進(jìn)入反向傳播引擎。
具體代碼如下：
// 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
};
// 對(duì)輸入的參數(shù)重新解析并賦值給新定義的變量tensors,grad_tensors等等，比如 inputs就是一個(gè)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; // 就是反向傳播的起點(diǎn)（根節(jié)點(diǎn)）
roots.reserve(num_tensors);
variable_list grads; // 就是反向傳播的梯度
grads.reserve(num_tensors);
// 依據(jù)輸入來(lái)配置roots和grads
for (int i = 0; i < num_tensors; i++) {
// tensors是輸入節(jié)點(diǎn)，即前向傳播圖的輸出
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增加一個(gè)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); // 增加一個(gè)梯度
}
}
// 構(gòu)建一個(gè)輸出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，用來(lái)判斷是否是葉子節(jié)點(diǎn)
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.
// 說(shuō)明是葉子節(jié)點(diǎn)
output_edges.emplace_back(std::make_shared<Identity>(), 0);
} else {
// 是中間節(jié)點(diǎn)
output_edges.emplace_back(grad_fn, output_nr);
}
}
}
// 現(xiàn)在，roots是包含有(前向傳播輸出節(jié)點(diǎn)的grad_fn_, 0)的vector。
// grads 是前向傳播產(chǎn)生的梯度，如果沒(méi)有配置，則初始化為(tensor(1.),)
// output_edges 是依據(jù)前向傳播輸入節(jié)點(diǎn) input 構(gòu)建的后向傳播輸出邊
variable_list outputs;
{
pybind11::gil_scoped_release no_gil;
auto& engine = python::PythonEngine::get_python_engine();
// 進(jìn)入引擎執(zhí)行
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
}
我們接下來(lái)分析 THPEngine_run_backward 用到的幾個(gè)輔助函數(shù)。
3.3.1 try_get_grad_accumulator
上面代碼之中，有 grad_fn = torch::autograd::impl::try_get_grad_accumulator(tensor) 來(lái)獲取計(jì)算梯度的方法。其實(shí)是用它來(lái)判斷是否是葉子節(jié)點(diǎn)，只有非葉子節(jié)點(diǎn) grad_accumulator_才不為空。
try_get_grad_accumulator 返回的是個(gè)指向Node對(duì)象的指針 : std::weak_ptr<Node> grad_accumulator_。就是如何計(jì)算梯度。
具體邏輯是：
- 先通過(guò)函數(shù)
get_autograd_meta返回一個(gè)AutogradMeta結(jié)構(gòu)體。 - 然后訪問(wèn)結(jié)構(gòu)體中的成員變量
grad_accumulator_，而grad_accumulator_是一個(gè)指向類型為Node對(duì)象的std::weak_ptr指針。 - 最后通過(guò)
lock()函數(shù)創(chuàng)建一個(gè)std::shared_ptr來(lái)管理對(duì)象。
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 被用來(lái)在輸入 tensor 基礎(chǔ)之上構(gòu)建一個(gè) Edge。
auto gradient_edge = torch::autograd::impl::gradient_edge(variable);
roots.push_back(std::move(gradient_edge)); // root增加一個(gè)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_ 之后，會(huì)賦予為 grad_fn，這樣就用來(lái)判斷是否為葉子節(jié)點(diǎn)。然后分別構(gòu)建葉子節(jié)點(diǎn)和中間節(jié)點(diǎn)，放到 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); // 葉子節(jié)點(diǎn)
} else {
output_edges.emplace_back(grad_fn, output_nr); // 非葉子節(jié)點(diǎn)
}
我們看看構(gòu)建output_edges的變量 grad_fn 和 output_nr，看看它們的由來(lái)。
grad_fn是通過(guò) try_get_grad_accumulator 方法得到的一個(gè)指向Node對(duì)象的std::shared_ptr指針，就是如何計(jì)算梯度的操作。
output_pr 由如下設(shè)置，其最終得到的是結(jié)構(gòu)體AutogradMeta中的成員變量uint32_t output_nr_。
const auto output_nr = tensor.output_nr();
emplace_back()函數(shù)向容器中中加入臨時(shí)對(duì)象， 臨時(shí)對(duì)象原地構(gòu)造，沒(méi)有賦值或移動(dòng)的操作。
回憶一下 Edge 的定義。所以可以看出來(lái)，emplace_back()就是使用了這些輸入生成了一個(gè) 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的第幾個(gè)輸入
};
輸入轉(zhuǎn)換如下圖，可以看出來(lái)輸入從 Python 如何進(jìn)行轉(zhuǎn)換最終傳入C++引擎，以如下變量為例：
- Python 的 tensors 被轉(zhuǎn)換為 C++ 的 root。
- Python 的 grad_tensors 被轉(zhuǎn)換為 C++ 的 grads。
- Python 的 inputs 被轉(zhuǎn)換為 C++ 的 output_edges。
- 最終把這三個(gè)變量傳遞給引擎：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 最終調(diào)用到了 PythonEngine 的處理邏輯。
auto& engine = python::PythonEngine::get_python_engine();
// 進(jìn)入引擎執(zhí)行
outputs = engine.execute(roots, grads, keep_graph, create_graph, accumulate_grad, output_edges);
3.4.1 獲取引擎
get_python_engine這里定義了一個(gè)靜態(tài)變量。整個(gè)PyTorch程序全局只維護(hù)一個(gè)Engine實(shí)例，也就是PythonEngine實(shí)例。
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 定義
所以我們來(lái)看看PythonEngine 定義。PythonEngine 是 Engine 的派生類，相當(dāng)于封裝了一下。主要是針對(duì)python世界的特點(diǎn)做了一些定制，比如：PythonEngine子類重寫(xiě)了父類的execute，把C++異常翻譯為Python異常的功能，核心工作還是由Engine基類來(lái)完成：
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 代碼如下，于是從下文開(kāi)始，我們要看看 Engine 是如何運(yùn)作的。
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)
手機(jī)如下：
3.5 另一調(diào)用途徑
最后，我們?cè)俨迦胍粋€(gè) run_backward 進(jìn)行分析。
run_backward 位于 torch/csrc/autograd/autograd.cpp。這里應(yīng)該是專門為了 C++ 世界直接調(diào)用的需要，與我們之前通過(guò) Python 迂回調(diào)用不同。
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 最后也調(diào)用了 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);
}
}
}
// 調(diào)用了引擎代碼
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;
}
至此，調(diào)用過(guò)程分析完畢，其核心就是調(diào)用引擎函數(shù)進(jìn)行處理，所以下一篇我們來(lái)開(kāi)始分析引擎。
浙公網(wǎng)安備 33010602011771號(hào)