[源碼解析]PyTorch如何實(shí)現(xiàn)前向傳播(1) --- 基礎(chǔ)類(上)
[源碼解析]PyTorch如何實(shí)現(xiàn)前向傳播(1) --- 基礎(chǔ)類(上)
0x00 摘要
本系列將通過大概十篇左右文章來分析 PyTorch 的自動微分功能如何實(shí)現(xiàn)。本文是前向傳播的第一篇，介紹自動微分（梯度計(jì)算）所涉及的部分 PyTorch 基礎(chǔ)類。因?yàn)樽謹(jǐn)?shù)太多（1萬兩千字），所以拆分成上下兩篇。
系列前幾篇連接如下：
深度學(xué)習(xí)利器之自動微分(3) --- 示例解讀
0x01 總體邏輯
為了行文完整，我們從前文結(jié)尾處摘取了總體邏輯關(guān)系如下。
如果從計(jì)算圖角度來看前向計(jì)算的過程，就是在構(gòu)建圖和執(zhí)行圖。"構(gòu)建圖"描述的是節(jié)點(diǎn)運(yùn)算之間的關(guān)系。"執(zhí)行圖"則是在會話中執(zhí)行這個(gè)運(yùn)算關(guān)系，就是張量在計(jì)算圖之中進(jìn)行前向傳播的過程。
前向計(jì)算依賴一些基礎(chǔ)類，在具體分析前向傳播之前，我們先要看看這些基礎(chǔ)類之間的邏輯關(guān)系。從DAG角度來分析 PyTorch 這個(gè)系統(tǒng)，其具體邏輯如下。
- 圖表示計(jì)算任務(wù)。PyTorch把計(jì)算都當(dāng)作是一種有向無環(huán)圖，或者說是計(jì)算圖，但這是一種虛擬的圖，代碼中沒有真實(shí)的數(shù)據(jù)結(jié)構(gòu)。
- 計(jì)算圖由節(jié)點(diǎn)（Node）和邊（Edge）組成。
- 節(jié)點(diǎn)（Node）代表了運(yùn)算操作。
- 一個(gè)節(jié)點(diǎn)通過邊來獲得 0 個(gè)或多個(gè)
Tensor，節(jié)點(diǎn)執(zhí)行計(jì)算之后會產(chǎn)生 0 個(gè)或多個(gè)Tensor。 - 節(jié)點(diǎn)的成員變量 next_functions 是一個(gè) tuple 列表，此列表就代表本節(jié)點(diǎn)要輸出到哪些其他 Function。列表個(gè)數(shù)就是這個(gè) grad_fn 的 Edge 數(shù)目，列表之中每一個(gè) tuple 對應(yīng)一條 Edge 信息，內(nèi)容就是 (Edge.function, Edge.input_nr)。
- 一個(gè)節(jié)點(diǎn)通過邊來獲得 0 個(gè)或多個(gè)
- 邊（Edge）就是運(yùn)算操作之間的流向關(guān)系。
- Edge.function ：表示此 Edge 需要輸出到哪一個(gè)其他 Function。
- Edge.input_nr ：指定本 Edge 是 Function 的第幾個(gè)輸入。
- 使用張量（ Tensor） 表示數(shù)據(jù)，就是在節(jié)點(diǎn)間流動的數(shù)據(jù)，如果沒有數(shù)據(jù)，計(jì)算圖就沒有任何意義。
具體可以參見下圖：
+---------------------+ +----------------------+
| SubBackward0 | | PowBackward0 |
| | Edge | | Edge
| next_functions +-----+--------> | next_functions +----------> ...
| | | | |
+---------------------+ | +----------------------+
|
|
| +----------------------+
| Edge | MulBackward0 |
+--------> | | Edge
| next_functions +----------> ...
| |
+----------------------+
0x02 廢棄類
我們先看看幾個(gè)已經(jīng)廢棄的類，這些類雖然廢棄了，但是代碼中依然有大量使用，網(wǎng)上也有大量文章與之相關(guān)，所以我們有必要先研究一下，我們在文章中可能會混用，還希望大家諒解。
2.1 Variable
早期版本之中，有Tensor和Variable兩種數(shù)據(jù)結(jié)構(gòu)來存儲數(shù)據(jù)，Tensor只負(fù)責(zé)多維數(shù)組的運(yùn)算。自動微分的職責(zé)是Variable完成的。Variable包含了與autograd有關(guān)的屬性，可以是計(jì)算圖中的葉子節(jié)點(diǎn)，也可以是計(jì)算時(shí)候產(chǎn)生的中間變量。
在0.4.0版本之后，Tensor和Variable 的功能進(jìn)行了合并，自動微分的使用就更加簡單了。現(xiàn)在，Variable 其實(shí)就是Tensor，只是為了向后兼容，才保留這個(gè)名字。
Variable (deprecated)
^^^^^^^^^^^^^^^^^^^^^
.. warning::
The Variable API has been deprecated: Variables are no longer necessary to
use autograd with tensors. Autograd automatically supports Tensors with
``requires_grad`` set to ``True``. Below please find a quick guide on what
has changed:
- ``Variable(tensor)`` and ``Variable(tensor, requires_grad)`` still work as expected,
but they return Tensors instead of Variables.
- ``var.data`` is the same thing as ``tensor.data``.
Variable 的定義在：torch/csrc/autograd/variable.h，我們可以看看注釋中 "Gradient Edges" 的相關(guān)部分。可以看出來，"Variable" 具有"gradient_edge"的概念，這是自動梯度計(jì)算圖的邊，在反向傳播之中用來把變量和梯度函數(shù)的特定輸入聯(lián)系起來。
更準(zhǔn)確地說，這個(gè)梯度函數(shù)可以是兩個(gè)函數(shù)之一：
grad_fn，如果variable 在圖的內(nèi)部。這是產(chǎn)生梯度變量的梯度函數(shù)。grad_accumulator，如果變量是一個(gè)葉子節(jié)點(diǎn)，它將一個(gè)標(biāo)量梯度值累加到它的'grad'變量之中。
namespace torch { namespace autograd {
/// `Variable` is exactly the same as `Tensor` (i.e. we have `using Variable = at::Tensor`).
/// This means you can perform all the usual mathematical and other
/// operations you can perform on `Tensor`s also on `Variable`s.
///
/// The only reason we are keeping the `Variable` class is backward compatibility
/// with external user's legacy C++ frontend code. Our intention is to eliminate
/// the `Variable` class in the near future.
using Variable = at::Tensor;
} // namespace autograd
} // namespace torch
/// Gradient Edges
///~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
/// Furthermore, `Variable`s have the notion of a `gradient_edge`, which is the
/// edge in the autograd graph that connects the variable to a particular input
/// of the gradient function that will be invoked with the variable during the
/// backward pass. More precisely, this gradient function can be one of two
/// things:
/// 1. A `grad_fn`, if the variable is in the interior of the graph. This is the
/// gradient of the function that produced the variable.
/// 2. A `grad_accumulator`, if the variable is a leaf, which accumulates a
/// scalar gradient value into its `grad` variable.
2.2 Function
我們結(jié)合前面 Variable 的概念來看，F(xiàn)unction 指的是在計(jì)算圖中某個(gè)節(jié)點(diǎn)所進(jìn)行的運(yùn)算，比如加減乘除卷積等等。每當(dāng)對Tensor施加一個(gè)運(yùn)算的時(shí)候，就會產(chǎn)生一個(gè)Function對象，它記錄運(yùn)算的輸入，記錄運(yùn)算的發(fā)生，產(chǎn)生運(yùn)算的結(jié)果。Tensor使用.grad_fn屬性記錄這個(gè)計(jì)算圖的入口。
Function 內(nèi)部有 forward() 和 backward() 兩個(gè)方法，分別應(yīng)用于正向、反向傳播。反向傳播過程中，autograd引擎會按照逆序，通過Function的backward依次計(jì)算梯度。
在最新的代碼中，Function 已經(jīng)被 Node 類替代，這樣是為了更好的表達(dá) 節(jié)點(diǎn) 這個(gè)概念。但是因?yàn)榕f代碼中依然使用了 Function，所以我們可能會混用這兩個(gè)概念。
Function 定義如下：
/// To use custom autograd operations, implement a Function subclass with
/// static forward and backward functions:
///
/// `forward` can take as many arguments as you want and should return either a
/// variable list or a Variable. Use of any direct Variable arguments will be
/// registered in the graph but no vectors/sets or any other data structures
/// will be traversed. You can use c10::optional<Tensor> as one of the arguments
/// and it will be registered as a variable in the graph if the argument has a
/// value. It should take a pointer to `torch::autograd::AutogradContext` as the
/// first argument. Variables can be saved in the `ctx` using
/// `ctx->save_for_backward`
/// (see `torch::autograd::AutogradContext::save_for_backward`) and other data
/// can be saved in the `ctx->saved_data` map
/// (see `torch::autograd::AutogradContext::saved_data`)
/// in the form of `<std::string, at::IValue>` pairs.
///
/// `backward` should take a pointer to `torch::autograd::AutogradContext`
/// and a variable list containing as many Variables as there were outputs from
/// `forward` as arguments. It should return as many Variables as there were
/// inputs with each of them containing the gradient w.r.t. its corresponding
/// input. Variables saved in `forward` can be accessed with
/// `ctx->get_saved_variables` (see
/// `torch::autograd::AutogradContext::get_saved_variables`) and other saved
/// data can be accessed from `ctx->saved_data`.
template <class T>
struct TORCH_API Function {
// We need to use a different template parameter than T here because T will
// inherit from Function, and when Function<T> is instantiated, T::forward
// is not declared yet.
// The enable_if check is to ensure that the user doesn't explicitly provide
// the parameter X.
template<typename X=T, typename... Args>
static auto apply(Args&&... args) -> std::enable_if_t<std::is_same<X,T>::value, forward_t<X,Args...>>;
};
0x03 Tensor
前面提到，計(jì)算圖構(gòu)成了前向/反向傳播的結(jié)構(gòu)基礎(chǔ)，而Tensor張量 是 PyTorch 中構(gòu)建計(jì)算圖的基礎(chǔ)之一。
Tensor是PyTorch實(shí)現(xiàn)多維數(shù)組計(jì)算和自動微分的關(guān)鍵數(shù)據(jù)結(jié)構(gòu)。
Tensor類似于numpy的ndarray，可以對Tensor進(jìn)行各種數(shù)學(xué)運(yùn)算；- 當(dāng)設(shè)置
.requires_grad = True，在Tensor之上進(jìn)行的各種操作就會被記錄下來，用于后續(xù)梯度計(jì)算。
3.1 定義 in python
我們看看第一個(gè)例子中運(yùn)行時(shí)的Tensor，其中可以看到Tensor的成員變量。
Q = {Tensor}
data = {Tensor} tensor(-12.)
device = {device} cpu
dtype = {dtype} torch.float32
grad = {NoneType} None
grad_fn = {SubBackward0}
metadata = {dict: 0} {}
next_functions = {tuple: 2}
0 = {tuple: 2} (<MulBackward0 object at 0x000001F9547A5848>, 0)
1 = {tuple: 2} (<PowBackward0 object at 0x000001F9547A53C8>, 0)
__len__ = {int} 2
requires_grad = {bool} True
is_cuda = {bool} False
is_leaf = {bool} False
is_meta = {bool} False
is_mkldnn = {bool} False
is_mlc = {bool} False
is_quantized = {bool} False
is_sparse = {bool} False
is_sparse_csr = {bool} False
is_vulkan = {bool} False
is_xpu = {bool} False
layout = {layout} torch.strided
name = {NoneType} None
names = {tuple: 0} ()
ndim = {int} 0
output_nr = {int} 0
requires_grad = {bool} True
shape = {Size: 0} torch.Size([])
我們看看其中的部分成員變量：
-
data：該張量的數(shù)據(jù)。
-
dtype ：該張量的數(shù)據(jù)類型。
-
device： 存放該張量的設(shè)備類型，比如 CPU 或者是 GPU。
-
grad：保存數(shù)據(jù)data對應(yīng)的梯度，和數(shù)據(jù)data的形狀一樣。
- PyTorch會自動追蹤和記錄對與張量的所有操作，當(dāng)前向計(jì)算完成后調(diào)用
.backward()方法會自動計(jì)算梯度并且將計(jì)算結(jié)果保存到grad屬性中。 - requires_grad = False時(shí)，grad為None。
- 梯度值不會自動清空，每次在backward計(jì)算時(shí)都需要將前一時(shí)刻的梯度歸零，否則梯度值會一直累加。
- PyTorch會自動追蹤和記錄對與張量的所有操作，當(dāng)前向計(jì)算完成后調(diào)用
-
grad_fn：指向一個(gè)Function對象。
- 這個(gè)Function對象用來在反向傳播時(shí)候計(jì)算輸入的梯度。
- 若本張量是非葉節(jié)點(diǎn)，則 Function 是向葉節(jié)點(diǎn)方向操作的反向傳播函數(shù)，比如例子里 O 節(jié)點(diǎn)對應(yīng)的函數(shù)就是MulBackward，即乘法操作的反向函數(shù)；
- 若本張量是葉節(jié)點(diǎn)且requires_grad為True，則 grad_fn 是None。
- grad_fn 有一個(gè)屬性 next_functions，這是一個(gè)二維 tuple，形式為( (函數(shù)1， 整數(shù)1)，(函數(shù)2，整數(shù)2), ..., (函數(shù)n，整數(shù)n) )。 后續(xù)我們會詳細(xì)解釋。
-
is_leaf：記錄該張量是否是葉子節(jié)點(diǎn) 。
- 用戶顯式初始化的張量是葉子節(jié)點(diǎn)。
- 所有requires_grad=False的張量按照慣例也是葉子節(jié)點(diǎn)。
is_leaf屬性只有在需要求導(dǎo)的時(shí)候才有意義。對于任意一個(gè)張量來說，我們可以用tensor.is_leaf來判斷它是否是葉子張量（leaf tensor）。在反向傳播過程中，只有is_leaf=True的時(shí)候，需要求導(dǎo)張量的導(dǎo)數(shù)結(jié)果才會被保留下來。- 對于葉子節(jié)點(diǎn)來說，其
grad_fn屬性都為空；而對于非葉子結(jié)點(diǎn)來說，因?yàn)樗鼈兪峭ㄟ^一些操作生成的，所以其grad_fn不為空。
-
requires_grad : 設(shè)置為
True則表示該Tensor需要求導(dǎo)，用于判斷該tensor是否需要被跟蹤并計(jì)算梯度。requires_grad屬性默認(rèn)為False，也就是Tensor變量默認(rèn)是不需要求導(dǎo)的。- 如果一個(gè)節(jié)點(diǎn)的
requires_grad是True，那么所有依賴它的節(jié)點(diǎn)的requires_grad也會是True。換言之，如果一個(gè)節(jié)點(diǎn)依賴的所有節(jié)點(diǎn)都不需要求導(dǎo)，那么它的requires_grad也會是False。因此在反向傳播過程中，該節(jié)點(diǎn)所在的子圖會被排除在計(jì)算過程之外。
Python的定義其實(shí)只是C++世界定義的一個(gè)映射，我們接下來就看看在C++如何定義。
3.2 查找定義
我們逐級找找 Tensor的定義。
首先來到：torch_C_VariableFunctions.pyi
def tensor(data: Any, dtype: Optional[_dtype]=None, device: Union[_device, str, None]=None, requires_grad: _bool=False) -> Tensor: ...
然后來到: torch/_tensor.py
3.2.1 Tensor
可以看到Tensor 的基類是 torch._C._TensorBase。
class Tensor(torch._C._TensorBase):
3.2.2 _TensorBase
_TensorBase 是動態(tài)生成的，代碼在比如python_stubs\xxx\torch\_C\_TensorBase.py
class _TensorBase(object):
我們在 torch/_C/__init__.pyi.in可以看到，torch._C._TensorBase 其實(shí)就是在 C++世界中定義的，但是需要導(dǎo)出到 python世界。
# Defined in torch/csrc/autograd/python_variable.cpp
class _TensorBase(metaclass=_TensorMeta):
requires_grad: _bool
shape: Size
data: Tensor
names: List[str]
device: _device
dtype: _dtype
layout: _layout
real: Tensor
imag: Tensor
T: Tensor
ndim: _int
output_nr: _int
_version: _int
_base: Optional[Tensor]
_cdata: _int
grad_fn: Any
_grad_fn: Any
_grad: Optional[Tensor]
_backward_hooks: Optional[Dict[_int, Callable[[Tensor], Optional[Tensor]]]]
${tensor_method_hints}
3.3 轉(zhuǎn)換
本文只是簡略看看如何從C++世界轉(zhuǎn)換到Python世界，在此處不做深入研究。
3.3.1 Python 導(dǎo)入
代碼中引入 PyTorch 是通過 import torch 完成的。Import torch 的時(shí)候，按照Python規(guī)范，位于torch/__init__.py中的邏輯就會被執(zhí)行，torch/__init__.py 的關(guān)鍵就是torch._C，代碼如下：
from torch._C import *
torch._C是C++編譯出來的共享庫文件，比如linux下的so文件。
Tensor類就是繼承自torch._C._TensorBase。導(dǎo)入了 torch._C就導(dǎo)入了torch._C._TensorBase，然后 torch.Tensor 就有了繼承的基礎(chǔ)。具體如下：
+---------------------------+
| import torch |
+------------+--------------+
|
|
v
+------------+--------------+
| torch/__init__.py |
| |
| from torch._C impor * |
| |
+------------+--------------+
|
|
v
+------------+--------------+
| torch._C._TensorBase |
+---------------------------+
所以我們接下來要看看 torch._C 是怎么來從 C++ 世界中導(dǎo)出到 python的。
3.3.2 C++ 導(dǎo)出 & 初始化
接下來我們看看C++世界如何導(dǎo)出了TensorBase。
要在Python中能夠import torch._C，則必須要使用Python的擴(kuò)展規(guī)范來導(dǎo)出這個(gè)符號。
3.3.2.1 共享庫入口
對于一個(gè) Python module，共享庫需要實(shí)現(xiàn) PyInit_modulename 符號來作為import時(shí)候的邏輯入口。對于PyTorch來說這個(gè)modulename 是_C。在torch/csrc/stub.cpp中 實(shí)現(xiàn)了PyInit__C這個(gè)函數(shù)。
#include <Python.h>
extern PyObject* initModule();
PyMODINIT_FUNC PyInit__C()
{
return initModule();
}
如果使用 JIT，則我們直接看 torch/csrc/deploy/interpreter/interpreter_impl.cpp，這里省略了眾多代碼。
struct ConcreteInterpreterImpl : public torch::deploy::InterpreterImpl {
ConcreteInterpreterImpl() {
PyImport_AppendInittab("torch._C", initModule);
}
這就是解釋器的代碼，里面也調(diào)用了 initModule。
3.3.2.2 initModule
initModule函數(shù)是對python環(huán)境中的torch module進(jìn)行初始化。其定義在 torch/csrc/Module.cpp，此處省略了眾多代碼。
PyObject* initModule() {
THPSize_init(module);
THPDtype_init(module);
THPDTypeInfo_init(module);
THPLayout_init(module);
THPMemoryFormat_init(module);
THPQScheme_init(module);
THPDevice_init(module);
THPStream_init(module);
ASSERT_TRUE(THPVariable_initModule(module)); // 繼續(xù)分析這里，其中會設(shè)定_TensorBase
ASSERT_TRUE(THPFunction_initModule(module));
ASSERT_TRUE(THPEngine_initModule(module));
}
initModule 調(diào)用 THPVariable_initModule，代碼在 torch/csrc/autograd/python_variable.cpp，這里會設(shè)定_TensorBase。
bool THPVariable_initModule(PyObject *module)
{
THPVariableMetaType.tp_base = &PyType_Type;
if (PyType_Ready(&THPVariableMetaType) < 0)
return false;
Py_INCREF(&THPVariableMetaType);
PyModule_AddObject(module, "_TensorMeta", (PyObject *)&THPVariableMetaType);
static std::vector<PyMethodDef> methods;
THPUtils_addPyMethodDefs(methods, torch::autograd::variable_methods);
THPUtils_addPyMethodDefs(methods, extra_methods);
THPVariableType.tp_methods = methods.data();
if (PyType_Ready(&THPVariableType) < 0)
return false;
Py_INCREF(&THPVariableType);
// 設(shè)定_TensorBase
PyModule_AddObject(module, "_TensorBase", (PyObject *)&THPVariableType);
torch::autograd::initTorchFunctions(module);
torch::autograd::initTensorImplConversion(module);
return true;
}
3.3.2.3 注冊TensorBase
執(zhí)行THPVariable_initModule的時(shí)候，使用如下代碼來將 THPVariableType 注冊成為torch._C._TensorBase。所以torch._C._TensorBase就是c++中的 THPVariableType。
PyModule_AddObject(module, "_TensorBase", (PyObject *)&THPVariableType);
我們來看看 THPVariableType。里面定義了很多函數(shù)。
PyTypeObject THPVariableType = {
PyVarObject_HEAD_INIT(&THPVariableMetaType, 0)
"torch._C._TensorBase", /* tp_name */
sizeof(THPVariable), /* tp_basicsize */
0, /* tp_itemsize */
(destructor)THPVariable_dealloc, /* tp_dealloc */
// 省略......
nullptr, /* tp_methods */
nullptr, /* tp_members */
THPVariable_properties, /* tp_getset */ // 重點(diǎn)在這里，注冊了函數(shù)
// 省略......
THPVariable_pynew, /* tp_new */
};
現(xiàn)在我們注冊了torch._C._TensorBase這個(gè)Python類，下面就要往這個(gè)類上注冊一些函數(shù)。
tp_getset 是Python虛擬機(jī)類機(jī)制里面的一個(gè)函數(shù)集，就是一個(gè) THPVariable_properties。以下是 _TenseBase 的函數(shù)集，我們可以看到 grad_fn 和 grad 這兩個(gè)熟悉的面孔。
static struct PyGetSetDef THPVariable_properties[] = {
{"T", (getter)THPVariable_get_T, nullptr, nullptr, nullptr},
{"_cdata", (getter)THPVariable_get_cdata, nullptr, nullptr, nullptr},
{"_version", (getter)THPVariable_get_version, nullptr, nullptr, nullptr},
{"grad_fn", (getter)THPVariable_get_grad_fn, nullptr, nullptr, nullptr},
{"_grad_fn", (getter)THPVariable_get_grad_fn, (setter)THPVariable_set_grad_fn, nullptr, nullptr},
{"is_leaf", (getter)THPVariable_is_leaf, nullptr, nullptr, nullptr},
{"data", (getter)THPVariable_get_data, (setter)THPVariable_set_data, nullptr, nullptr},
{"_grad", (getter)THPVariable_get_grad, (setter)THPVariable_set_grad, nullptr, nullptr}, // Allows the python class to override .grad
{"grad", (getter)THPVariable_get_grad, (setter)THPVariable_set_grad, nullptr, nullptr},
{"_base", (getter)THPVariable_get_base, nullptr, nullptr, nullptr},
{"volatile", (getter)THPVariable_get_volatile, (setter)THPVariable_set_volatile, nullptr, nullptr},
{"output_nr", (getter)THPVariable_get_output_nr, nullptr, nullptr, nullptr},
{"requires_grad", (getter)THPVariable_get_requires_grad, (setter)THPVariable_set_requires_grad, nullptr, nullptr},
{"_backward_hooks", (getter)THPVariable_get_backwards_hooks, (setter)THPVariable_set_backwards_hooks, nullptr, nullptr},
{"name", (getter)THPVariable_get_name, nullptr, nullptr, nullptr},
{"shape", (getter)THPVariable_get_shape, nullptr, nullptr, nullptr},
{"is_cuda", (getter)THPVariable_is_cuda, nullptr, nullptr, nullptr},
{"is_xpu", (getter)THPVariable_is_xpu, nullptr, nullptr, nullptr},
{"is_sparse", (getter)THPVariable_is_sparse, nullptr, nullptr, nullptr},
{"is_sparse_csr", (getter)THPVariable_is_sparse_csr, nullptr, nullptr, nullptr},
{"is_mkldnn", (getter)THPVariable_is_mkldnn, nullptr, nullptr, nullptr},
{"is_mlc", (getter)THPVariable_is_mlc, nullptr, nullptr, nullptr},
{"is_vulkan", (getter)THPVariable_is_vulkan, nullptr, nullptr, nullptr},
{"is_complex", (getter)THPVariable_is_complex, nullptr, nullptr, nullptr},
{"is_quantized", (getter)THPVariable_is_quantized, nullptr, nullptr, nullptr},
{"is_meta", (getter)THPVariable_is_meta, nullptr, nullptr, nullptr},
{"dtype", (getter)THPVariable_dtype, nullptr, nullptr, nullptr},
{"layout", (getter)THPVariable_layout, nullptr, nullptr, nullptr},
{"device", (getter)THPVariable_device, nullptr, nullptr, nullptr},
{"ndim", (getter)THPVariable_get_ndim, nullptr, nullptr, nullptr},
{"names", (getter)THPVariable_get_names, (setter)THPVariable_set_names, nullptr, nullptr},
{"real", (getter)THPVariable_get_real, (setter)THPVariable_set_real, nullptr, nullptr},
{"imag", (getter)THPVariable_get_imag, (setter)THPVariable_set_imag, nullptr, nullptr},
{nullptr}
};
這個(gè)初始化邏輯和映射邏輯如下：
Python + C++ +---------------+
| | |
+---------------------------+ | | PyInit__C |
| import torch | | | |
+------------+--------------+ | +-------+-------+
| | |
| | |
v | |
+------------+--------------+ | v
| torch/__init__.py | | +-------+-------+
| | | | initModule |
| from torch._C impor * | | +-------+-------+
| | | |
+------------+--------------+ | |
| | |
| | v
| | +--------------+----------------+
| | | |
| | | THPVariable_initModule(module)|
| | | |
| | +--------------+----------------+
| | |
| | |
| | |
| | v
| | +-------------------------------+---------------------------------------+
| | | |
| | | PyModule_AddObject(module, "_TensorBase",(PyObject *)&THPVariableType)|
| | | |
| | +-------------------------------+---------------------------------------+
| | |
| | |
| | |
| | v
| | +-----------+--------------+ +------------------------------------------------------+
| | | THPVariableType | | THPVariable_properties+ |
v | | | | |
+------------+--------------+ | | | | |
| torch._C._TensorBase | <-----------------------> | tp_getset -----> | { grad, grad_fn, T, _cdata, is_leaf, output_nr ...} |
+---------------------------+ | | | | |
| +--------------------------+ +------------------------------------------------------+
+
手機(jī)如下：
3.4 next_functions 設(shè)置
因?yàn)?next_functions 是精髓，而 next_functions 是在 autograd 之中設(shè)置，于是我們需要看看初始化autograd 過程。然后才能知道如何設(shè)置 next_functions。
3.5 初始化autograd
我們以 AccumulateGrad 為例來看看如何初始化。
首先看看 AccumulateGrad 的定義，這里省略了 AccumulateGrad 部分成員函數(shù)。從構(gòu)建函數(shù)可看出來，一個(gè)AccumulateGrad實(shí)例必須用一個(gè)Variable構(gòu)建，內(nèi)部成員變量就是Variable variable。apply調(diào)用接收一個(gè)Variable list 實(shí)例，這和Variable grad_accumulator_相關(guān)。
struct TORCH_API AccumulateGrad : public Node {
explicit AccumulateGrad(Variable variable_);
variable_list apply(variable_list&& grads) override;
Variable variable;
};
舊版本之中，定義如下：
struct AccumulateGrad : public Function {
explicit AccumulateGrad(Variable variable_);
variable_list apply(variable_list&& grads) override;
Variable variable;
};
接下來看看如何初始化 AccumulateGrad。
3.5.1 擴(kuò)展
在initModule()函數(shù)初始化完畢之后，import torch 的初始化工作還沒有結(jié)束。python的初始化腳本還要繼續(xù)處理很多模塊，比如torch/__init__.py 文件中有：
# Check to see if we can load C extensions, and if not provide some guidance
# on what the problem might be.
try:
# _initExtension is chosen (arbitrarily) as a sentinel.
from torch._C import _initExtension
_initExtension 會調(diào)用到 _C._initExtension(manager_path())。_C._initExtension對應(yīng)的是 THPModule_initExtension。
static PyMethodDef TorchMethods[] = {
{"_initExtension", THPModule_initExtension, METH_O, nullptr},
// ....
}
THPModule_initExtension 函數(shù)會調(diào)用THPAutograd_initFunctions，該方法初始化了自動微分系統(tǒng)。
// Callback for python part. Used for additional initialization of python classes
static PyObject * THPModule_initExtension(PyObject *_unused, PyObject *shm_manager_path)
{
// 省略代碼
THPQInt8Storage_postInit(module);
THPQInt32Storage_postInit(module);
THPBFloat16Storage_postInit(module);
THPComplexDoubleStorage_postInit(module);
THPComplexFloatStorage_postInit(module);
THPAutograd_initFunctions(); // 這里調(diào)用,初始化了微分系統(tǒng)
// 省略代碼
}
THPAutograd_initFunctions 就是在 _TensorBase 基礎(chǔ)之上，再加入新的屬性或者函數(shù)集。**這里會調(diào)用了addClass 方法，把 AccumulateGrad 和 accumulate_grad_properties 聯(lián)系在一起 **。
void THPAutograd_initFunctions()
{
THPObjectPtr module(PyModule_New("torch._C._functions"));
if (!module) throw python_error();
static PyTypeObject AccumulateGradClass;
addClass<AccumulateGrad, NoCtor>(module, AccumulateGradClass, "AccumulateGrad", accumulate_grad_properties); // AccumulateGrad 相關(guān)
static PyTypeObject CopyBackwardsClass;
addClass<CopyBackwards, NoCtor>(module, CopyBackwardsClass, "CopyBackwards");
// 省略其他
}
3.5.2 addClass
addClass 會調(diào)用到 registerCppFunction 注冊 type（ function_properties），我們這里參數(shù) function_properties 就是 accumulate_grad_properties，type 就是 AccumulateGradClass。
template<typename C, typename T>
static void addClass(PyObject* module, PyTypeObject& type, const char* name,
PyGetSetDef* function_properties=nullptr, PyMethodDef* function_methods=nullptr)
{
// 這里設(shè)置了 accumulate_grad_properties
createForwardFunctionPyTypeObject<T>(type, name, function_properties, function_methods);
Py_INCREF(&type);
PyModule_AddObject(module, name, (PyObject*)&type);
// 注冊了 type
registerCppFunction(typeid(C), &type);
}
這里有兩組操作，一個(gè)是 createForwardFunctionPyTypeObject，一個(gè)是 registerCppFunction。我們逐一看看。我們先看 registerCppFunction，然后看 createForwardFunctionPyTypeObject。
3.5.2.1 accumulate_grad_properties
前面提到，addClass 方法，把 AccumulateGrad 和 accumulate_grad_properties 聯(lián)系在一起。具體來說，就是通過 createForwardFunctionPyTypeObject 把 accumulate_grad_properties 聯(lián)系起來。
accumulate_grad_properties 定義在 torch/csrc/autograd/functions/init.cpp
static struct PyGetSetDef accumulate_grad_properties[] = {
THP_FUNCTION_DEFAULT_PROPERTIES,
{(char*)"variable", accumulateGradVar, nullptr, nullptr, nullptr},
{nullptr}
};
THP_FUNCTION_DEFAULT_PROPERTIES 的定義在 torch/csrc/autograd/python_cpp_function.h
#define THP_FUNCTION_DEFAULT_PROPERTIES \
{(char*)"next_functions", (getter)THPCppFunction_next_functions, nullptr, nullptr, nullptr}, \
{(char*)"requires_grad", (getter)THPCppFunction_requires_grad, nullptr, nullptr, nullptr}, \
{(char*)"metadata", (getter)THPCppFunction_metadata, nullptr, nullptr, nullptr}
PyObject* THPCppFunction_next_functions(THPCppFunction* self, PyObject* hook);
PyObject* THPCppFunction_metadata(THPCppFunction *self, void *_unused);
PyObject* THPCppFunction_requires_grad(THPCppFunction* self, void *_unused);
所以，accumulate_grad_properties 就是拓展了 THP_FUNCTION_DEFAULT_PROPERTIES 和 accumulateGradVar。
static struct PyGetSetDef accumulate_grad_properties[] = {
// 這里是我們關(guān)注的
{(char*)"next_functions", (getter)THPCppFunction_next_functions, nullptr, nullptr, nullptr},
{(char*)"requires_grad", (getter)THPCppFunction_requires_grad, nullptr, nullptr, nullptr},
{(char*)"metadata", (getter)THPCppFunction_metadata, nullptr, nullptr, nullptr}
{(char*)"variable", accumulateGradVar, nullptr, nullptr, nullptr},
{nullptr}
};
具體邏輯如下，這里面就有 THPCppFunction_next_functions:
+-----------------------------------------------------------------------+
|accumulate_grad_properties |
| |
| |
| |
| "variable", accumulateGradVar |
| |
| |
| "next_functions", (getter)THPCppFunction_next_functions |
| |
| |
| "requires_grad", (getter)THPCppFunction_requires_grad |
| |
| |
| "metadata", (getter)THPCppFunction_metadata |
| |
+-----------------------------------------------------------------------+
3.5.2.3 createForwardFunctionPyTypeObject
createForwardFunctionPyTypeObject 是用來設(shè)置accumulate_grad_properties，具體函數(shù)如下：
template<typename Ctor>
PyTypeObject* createForwardFunctionPyTypeObject(PyTypeObject& type, const char* name,
PyGetSetDef* function_properties=nullptr, PyMethodDef* function_methods=nullptr)
{
type.tp_new = &CppFunction_pynew<Ctor>;
return _initFunctionPyTypeObject(type, name, function_properties, function_methods);
}
_initFunctionPyTypeObject 就是把 function_properties 設(shè)置到 tp_getset 之上。
PyTypeObject* _initFunctionPyTypeObject(PyTypeObject& type, const char* name,
PyGetSetDef* function_properties, PyMethodDef* function_methods)
{
type.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC;
type.tp_name = name;
type.tp_basicsize = sizeof(THPCppFunction);
type.tp_call = THPCppFunction_call;
type.tp_methods = function_methods ? function_methods : default_methods;
// 這里把 function_properties 設(shè)置到 tp_getset 之上
type.tp_getset = function_properties ? function_properties : default_properties;
type.tp_dealloc = THPCppFunction_dealloc;
type.tp_traverse = THPCppFunction_traverse;
type.tp_clear = THPCppFunction_clear;
if (PyType_Ready(&type) < 0) {
auto msg = std::string("Unable to instantiate PyTypeObject for ") + name;
throw std::runtime_error(msg);
}
return &type;
}
所以就把 THPCppFunction_next_functions 添加到了 AccumulateGradClass 的 next_functions 之上。即 AccumulateGradClass 有一個(gè)函數(shù)集，其中 next_functions 對應(yīng)了 THPCppFunction_next_functions。
+---------------------+
| AccumulateGradClass |
| |
| tp_getset |
| + |
| | |
+---------------------+
|
|
v
+-----------+-----------------------------------------------------------+
|accumulate_grad_properties |
| |
| |
| |
| "variable", accumulateGradVar |
| |
| |
| "next_functions", (getter)THPCppFunction_next_functions |
| |
| |
| "requires_grad", (getter)THPCppFunction_requires_grad |
| |
| |
| "metadata", (getter)THPCppFunction_metadata |
| |
+-----------------------------------------------------------------------+
我們回憶一下前面提到的 _TenseBase 來對比：
tp_getset 是Python虛擬機(jī)類機(jī)制里面的一個(gè)函數(shù)集，就是一個(gè) THPVariable_properties。
- AccumulateGradClass 設(shè)置的是 accumulate_grad_properties。
- _TenseBase 設(shè)置的是 THPVariable_properties。
以下是 _TenseBase 的函數(shù)集（我們省略了很多）。
static struct PyGetSetDef THPVariable_properties[] = {
{"grad_fn", (getter)THPVariable_get_grad_fn, nullptr, nullptr, nullptr},
{"_grad_fn", (getter)THPVariable_get_grad_fn, (setter)THPVariable_set_grad_fn, nullptr, nullptr},
{"is_leaf", (getter)THPVariable_is_leaf, nullptr, nullptr, nullptr},
{"data", (getter)THPVariable_get_data, (setter)THPVariable_set_data, nullptr, nullptr},
{"_grad", (getter)THPVariable_get_grad, (setter)THPVariable_set_grad, nullptr, nullptr}, // Allows the python class to override .grad
{"grad", (getter)THPVariable_get_grad, (setter)THPVariable_set_grad, nullptr, nullptr},
{"_base", (getter)THPVariable_get_base, nullptr, nullptr, nullptr},
{"output_nr", (getter)THPVariable_get_output_nr, nullptr, nullptr, nullptr},
{"requires_grad", (getter)THPVariable_get_requires_grad, (setter)THPVariable_set_requires_grad, nullptr, nullptr},
{"_backward_hooks", (getter)THPVariable_get_backwards_hooks,
.....
};
至此，業(yè)務(wù)邏輯如下：
Python + C++
|
+--------------------------------------+ | +---------------------------+
| torch/__init__.py | | | |
| | | | THPModule_initExtension |
| from torch._C import _initExtension | | | |
| | | +--------------+------------+
+-------------------+------------------+ | |
| | |
| | v
| | +---------------+--------------+
| | | |
| | | THPAutograd_initFunctions() |
| | | |
| | +---------------+--------------+
| | |
| | |
| | v
| | +---------------+-------------------------------------------+
| | | |
| | | addClass<AccumulateGrad, NoCtor>(module, |
| import | | AccumulateGradClass, |
| | | "AccumulateGrad", |
| | | accumulate_grad_properties) |
| | | |
| | +--------------+--------------------------------------------+
| | |
| | | register
v | v
| +----------------------------------------------------------+
+----------------------+ | +--------------------+ +---------------------+ |accumulate_grad_properties |
| | | | | | AccumulateGradClass | | |
| AccumulateGrad | <------------> | AccumulateGrad +-----> | | | "variable", accumulateGradVar |
| | | | | | tp_getset +-------> | |
| | | | | | | | "next_functions", (getter)THPCppFunction_next_functions |
+----------------------+ | +--------------------+ | | | |
| +---------------------+ | "requires_grad", (getter)THPCppFunction_requires_grad |
| | |
| | "metadata", (getter)THPCppFunction_metadata |
| | |
| +----------------------------------------------------------+
手機(jī)如下：
3.5.2.4 next_functions
THPCppFunction_next_functions 定義在 torch/csrc/autograd/python_cpp_function.cpp，其就是遍歷 next_edges_，然后提取出一個(gè)tuple列表，每個(gè)tuple 內(nèi)容就是 (Edge.function, Edge.input_nr)，最后作為 next_functions 進(jìn)行返回。
PyObject* THPCppFunction_next_functions(THPCppFunction* self, PyObject* hook)
{
const auto num_next = self->cdata->num_outputs();
THPObjectPtr py_functions(PyTuple_New(num_next));
if (!py_functions) return nullptr;
for (size_t i = 0; i < num_next; ++i) { // 遍歷
auto& c_tuple = self->cdata->next_edge(i); // 獲取 Edge
THPObjectPtr tuple(PyTuple_New(2));
if (!tuple) return nullptr;
PyObject *py_fn = functionToPyObject(c_tuple.function); // py_fn 就是 Edge.function
if (!py_fn) return nullptr;
PyTuple_SET_ITEM(tuple.get(), 0, py_fn);
PyObject *py_idx = THPUtils_packUInt32(c_tuple.input_nr); // py_idx 就是 Edge.input_nr
if (!py_idx) return nullptr;
PyTuple_SET_ITEM(tuple.get(), 1, py_idx);
// tuple 就是 (py_fn, py_idx)，就是 (Edge.function, Edge.input_nr)
PyTuple_SET_ITEM(py_functions.get(), i, tuple.release()); // 設(shè)置 py_functions的第幾個(gè)item
}
return py_functions.release(); // 返回tuple
}
next_edge 定義在 torch/csrc/autograd/function.h，其是 Node 的成員函數(shù)，而 返回的是 Edge 列表，而 AccumulateGrad 就是 Node 的派生類。
struct TORCH_API Node : std::enable_shared_from_this<Node> {
const Edge& next_edge(size_t index) const noexcept {
return next_edges_[index];
}
edge_list next_edges_; // 前向過程中的輸入variable，在前向過程中與該算子相關(guān)聯(lián)的邊
}
Edge 定義如下：
struct Edge {
/// The function this `Edge` points to.
std::shared_ptr<Node> function; // 指向目標(biāo)的Node
/// The identifier of a particular input to the function.
uint32_t input_nr; //指定本Edge是function的第幾個(gè)輸入
};
3.5.3 next_functions 性質(zhì)
所以我們以 AccumulateGrad 為例總結(jié)以下。
- grad_fn 有一個(gè)屬性 next_functions ，這是一個(gè)二維的tuple，形式為( (函數(shù)1， 整數(shù)1)，(函數(shù)2，整數(shù)2), ..., (函數(shù)N，整數(shù)N) )。
- next_functions 是一個(gè) tuple 列表，列表個(gè)數(shù)就是這個(gè) grad_fn 的 Edge 數(shù)目，列表之中每一個(gè) tuple 對應(yīng)一條 Edge 信息，內(nèi)容就是 (Edge.function, Edge.input_nr)。這個(gè)列表是由 THPCppFunction_next_functions 生成的。
- AccumulateGrad 的 next_functions 指向的就是一個(gè) tuple 列表（就是下圖中的 2），這個(gè)列表來自 AccumulateGradClass（就是下圖中的 1）。反向傳播時(shí)候，順著這個(gè) next_functions 就可以逐次計(jì)算梯度。
大致如下：
+-----------------+ +-----------------------+ +----------------------+ +---------------------+
| Tensor | | SubBackward0 | | PowBackward0 | | AccumulateGrad |
| | | | | | | |
| grad_fn +---->+ next_functions +-----+--> | next_functions +----> | next_functions +----> {}
| | | | | | | | |
+-----------------+ +-----------------------+ | +----------------------+ +---------------------+
|
|
| +----------------------+ +----------------------+ +---------------------+
| | MulBackward0 | | PermuteBackward | | AccumulateGrad |
+--> | | | | | |
| next_functions +----> | next_functions +----> | next_functions +-----+
| | | | | | |
+---------------------+ ++-------------------- -+ +----------------------+ +---------------------+ |
| AccumulateGradClass | |
| | |
| tp_getset | 2. point to the tuple list
| + | |
| | | |
+---------------------+ |
| v
|
v +-----> { (function 1, int 1), (function 2, int 2) ... (function n, int n) }
+-----------+-----------------------------------------------------+ |
|accumulate_grad_properties | |
| | |
| "variable", accumulateGradVar | |
| | |
| "next_functions", (getter)THPCppFunction_next_functions +--------+
| | 1. generate the tuple list
| "requires_grad", (getter)THPCppFunction_requires_grad |
| |
| "metadata", (getter)THPCppFunction_metadata |
| |
+-----------------------------------------------------------------+
手機(jī)如下：
至此，部分基礎(chǔ)類解析完畢，因?yàn)槲淖炙蓿覀儗⒃谙乱黄^續(xù)分析其他基礎(chǔ)類。
0xFF 參考
https://github.com/KeithYin/read-pytorch-source-code/
pytorch學(xué)習(xí)筆記（十三）：backward過程的底層實(shí)現(xiàn)解析
pytorch的自動求導(dǎo)機(jī)制 - 計(jì)算圖的建立
How autograd encodes the history
https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html
pytorch筆記(計(jì)算圖+autograd)-Node(1)
詳解Pytorch中的網(wǎng)絡(luò)構(gòu)造
計(jì)算圖——用Pytorch解釋李宏毅老師PPT中的實(shí)例
PyTorch自動求導(dǎo)（Autograd）原理解析
pytorch自動求導(dǎo)Autograd系列教程（一）
浙公網(wǎng)安備 33010602011771號