PyTorch內部機制解析：如何通過PyTorch實現Tensor
PyTorch 中的基本單位是張量(Tensor)。本文的主旨是如何在 PyTorch 中實現 Tensor 的概述，以便用戶可從 Python shell 與之交互。本文主要回答以下四個主要問題：
- PyTorch 如何通過擴展 Python 解釋器來定義可以從 Python 代碼中調用的 Tensor 類型?
- PyTorch 如何封裝實際定義 Tensor 屬性和方法的 C 的類庫?
- PyTorch 的 C 類包裝器如何生成 Tensor 方法的代碼?
- PyTorch 的編譯系統如何編譯這些組件并生成可運行的應用程序?
擴展 Python 解釋器
PyTorch 定義了一個新的包 torch。本文中，我們將考慮._C 模塊。這是一個用 C 編寫的被稱為「擴展模塊」的 Python 模塊，它允許我們定義新的內置對象類型(例如 Tensor)和調用 C / C ++函數。
._C 模塊定義在 torch/csrc/Module.cpp 文件中。init_C()/ PyInit__C()函數創建模塊并根據需要添加方法定義。這個模塊被傳遞給一些不同的__init()函數，這些函數會添加更多的對象到模塊中，以及注冊新的類型等。
__init() 可調用的部分函數如下：
- ASSERT_TRUE(THPDoubleTensor_init(module));
- ASSERT_TRUE(THPFloatTensor_init(module));
- ASSERT_TRUE(THPHalfTensor_init(module));
- ASSERT_TRUE(THPLongTensor_init(module));
- ASSERT_TRUE(THPIntTensor_init(module));
- ASSERT_TRUE(THPShortTensor_init(module));
- ASSERT_TRUE(THPCharTensor_init(module));
- ASSERT_TRUE(THPByteTensor_init(module));
這些__init()函數將每種類型的 Tensor 對象添加到._C 模塊，以便它們可以在._C 模塊中調用。下面我們來了解這些方法的工作原理。
THPTensor 類型
PyTorch 很像底層的 TH 和 THC 類庫，它定義了一個專門針對多種不同的類型數據的「通用」Tensor。在考慮這種專業化的工作原理之前，我們首先考慮如何在 Python 中定義新的類型，以及如何創建通用的 THPTensor 類型。
Python 運行時將所有 Python 對象都視為 PyObject * 類型的變量，PyObject * 是所有 Python 對象的「基本類型」。每個 Python 類型包含對象的引用計數，以及指向對象的「類型對象」的指針。類型對象確定類型的屬性。例如，該對象可能包含一系列與類型相關聯的方法，以及調用哪些 C 函數來實現這些方法。該對象還可能包含表示其狀態所需的任意字段。
定義新類型的準則如下：
- 創建一個結構體，它定義了新對象將包括的屬性
- 定義類型的類型對象
結構體本身可能十分簡單。在 Python 中，實際上所有浮點數類型都是堆上的對象。Python float 結構體定義為：
- typedef struct {
- PyObject_HEAD
- double ob_fval;
- } PyFloatObject;
PyObject_HEAD 是引入實現對象的引用計數的代碼的宏，以及指向相應類型對象的指針。所以在這種情況下，要實現浮點數，所需的唯一其他「狀態」是浮點值本身。
現在，我們來看看 THPTensor 類型的結構題：
- struct THPTensor {
- PyObject_HEAD
- THTensor *cdata;
- };
很簡單吧?我們只是通過存儲一個指針來包裝底層 TH 張量。關鍵部分是為新類型定義「類型對象」。我們的 Python 浮點數的類型對象的示例定義的形式如下：
- static PyTypeObject py_FloatType = {
- PyVarObject_HEAD_INIT(NULL, 0)
- "py.FloatObject", /* tp_name */
- sizeof(PyFloatObject), /* tp_basicsize */
- 0, /* tp_itemsize */
- 0, /* tp_dealloc */
- 0, /* tp_print */
- 0, /* tp_getattr */
- 0, /* tp_setattr */
- 0, /* tp_as_async */
- 0, /* tp_repr */
- 0, /* tp_as_number */
- 0, /* tp_as_sequence */
- 0, /* tp_as_mapping */
- 0, /* tp_hash */
- 0, /* tp_call */
- 0, /* tp_str */
- 0, /* tp_getattro */
- 0, /* tp_setattro */
- 0, /* tp_as_buffer */
- Py_TPFLAGS_DEFAULT, /* tp_flags */
- "A floating point number", /* tp_doc */
- };
想象一個類型對象的最簡單的方法就是定義一組該對象屬性的字段。例如，tp_basicsize 字段設置為 sizeof(PyFloatObject)。這是為了讓 Python 知道 PyFloatObject 調用 PyObject_New()時需要分配多少內存。你可以設置的字段的完整列表在 CPython 后端的 object.h 中定義：
https://github.com/python/cpython/blob/master/Include/object.h.
THPTensor 的類型對象是 THPTensorType，它定義在 csrc/generic/Tensor.cpp 文件中。該對象定義了 THPTensor 的類型名稱、大小及映射方法等。
我們來看看我們在 PyTypeObject 中設置的 tp_new 函數：
- PyTypeObject THPTensorType = {
- PyVarObject_HEAD_INIT(NULL, 0)
- ...
- THPTensor_(pynew), /* tp_new */
- };
tp_new 函數可以創建對象。它負責創建(而不是初始化)該類型的對象，相當于 Python 中的__new()__方法。C 實現是一個靜態方法，該方法傳遞實例化的類型和任意參數，并返回一個新創建的對象。
- static PyObject * THPTensor_(pynew)(PyTypeObject *type, PyObject *args, PyObject *kwargs)
- {
- HANDLE_TH_ERRORS
- Py_ssize_t num_args = args ? PyTuple_Size(args) : 0;
- THPTensorPtr self = (THPTensor *)type->tp_alloc(type, 0);
- // more code below
我們的新函數的***件事就是為 THPTensor 分配內存。然后，它會根據傳遞給該函數的參數進行一系列的初始化。例如，當從另一個 THPTensor y 創建 THPTensor x 時，我們將新創建的 THPTensor 的 cdata 字段值設置為以 y 的底層 TH Tensor 作為參數并調用 THTensor_(newWithTensor)返回的結果。這一過程中有內存大小、存儲、NumPy 數組和序列的類似的構造函數。
注意，我們只使用了 tp_new 函數，而不是同時使用 tp_new 和 tp_init(對應于 Python 中的 __init()__函數)。
Tensor.cpp 中定義的另一個重要的部分是索引的工作原理。PyTorch Tensors 支持 Python 的映射協議。這樣我們可以做如下事情：
- x = torch.Tensor(10).fill_(1)
- y = x[3] // y == 1
- x[4] = 2
- // etc.
注意，此索引可以拓展到多維 Tensor。
我們可以通過定義
https://docs.python.org/3.7/c-api/typeobj.html#c.PyMappingMethods 里描述的三種映射方法來使用[]符號。
最重要的方法是 THPTensor_(getValue)和 THPTensor_(setValue)，它們解釋了如何對 Tensor 進行索引，并返回一個新的 Tensor / Scalar(標量)，或更新現有 Tensor 的值。閱讀這些實現代碼，以更好地了解 PyTorch 是如何支持基本張量索引的。
通用構建(***部分)
我們可以花費大量時間探索 THPTensor 的各個方面，以及它如何與一個新定義 Python 對象相關聯。但是我們仍然需要明白 THPTensor_(init)()函數是如何轉換成我們在模塊初始化中使用的 THPIntTensor_init()函數。我們該如何使用定義「通用」Tensor 的 Tensor.cpp 文件，并使用它來生成所有類型序列的 Python 對象?換句話說，Tensor.cpp 里遍布著如下代碼：
- return THPTensor_(New)(THTensor_(new)(LIBRARY_STATE_NOARGS));
這說明了我們需要使類型特定的兩種情況：
- 我們的輸出代碼將調用 THP
Tensor_New(...)代替調用 THPTensor_(New) - 我們的輸出代碼將調用 TH
Tensor_new(...)代替調用 THTensor_(new)
換句話說，對于所有支持的 Tensor 類型，我們需要「生成」已經完成上述替換的源代碼。這是 PyTorch 的「構建」過程的一部分。PyTorch 依賴于配置工具(https://setuptools.readthedocs.io/en/latest/)來構建軟件包，我們在頂層目錄中定義一個 setup.py 文件來自定義構建過程。
使用配置工具構建擴展模塊的一個組件是列出編譯中涉及的源文件。但是，我們的 csrc/generic/Tensor.cpp 文件未列出!那么這個文件中的代碼最終是如何成為最終產品的一部分呢?
回想前文所述，我們從以上的 generic 目錄中調用 THPTensor *函數(如 init)。如果我們來看一下這個目錄，會發現一個定義了的 Tensor.cpp 文件。此文件的***一行很重要：
- //generic_include TH torch/csrc/generic/Tensor.cpp
請注意，雖然這個 Tensor.cpp 文件被 setup.py 文件引用，但它被包裝在一個叫 Python helper 的名為 split_types 的函數里。這個函數需要輸入一個文件，并在該文件內容中尋找「//generic_include」字符串。如果能匹配該字符串，它將會為每個張量類型生成一個具有以下變動的輸出文件，：
1. 輸出文件重命名為 Tensor
2. 輸出文件小幅修改如下：
- // Before:
- //generic_include TH torch/csrc/generic/Tensor.cpp
- // After:
- #define TH_GENERIC_FILE "torch/src/generic/Tensor.cpp"
- #include "TH/THGenerate<Type>Type.h"
引入第二行的頭文件有些許弊端，例如，引入了一些額外的上下文中定義的 Tensor.cpp 源代碼。讓我們看看其中一個頭文件：
- #ifndef TH_GENERIC_FILE
- #error "You must define TH_GENERIC_FILE before including THGenerateFloatType.h"
- #endif
- #define real float
- #define accreal double
- #define TH_CONVERT_REAL_TO_ACCREAL(_val) (accreal)(_val)
- #define TH_CONVERT_ACCREAL_TO_REAL(_val) (real)(_val)
- #define Real Float
- #define THInf FLT_MAX
- #define TH_REAL_IS_FLOAT
- #line 1 TH_GENERIC_FILE
- #include TH_GENERIC_FILE
- #undef accreal
- #undef real
- #undef Real
- #undef THInf
- #undef TH_REAL_IS_FLOAT
- #undef TH_CONVERT_REAL_TO_ACCREAL
- #undef TH_CONVERT_ACCREAL_TO_REAL
- #ifndef THGenerateManyTypes
- #undef TH_GENERIC_FILE
- #endif
這樣做的目的是從通用 Tensor.cpp 文件引入代碼，并使用后面的宏定義。例如，我們將 real 定義為一個浮點數，所以泛型 Tensor 實現中的任何代碼將指向一個 real 對象，實際上 real 被替換為浮點數。在對應的文件 THGenerateIntType.h 中，同樣的宏定義將用 int 替換 real。
這些輸出文件從 split_types 返回，并添加到源文件列表中，因此我們可以看到不同的類型的.cpp 代碼是如何創建的。
這里需要注意以下幾點：***，split_types 函數不是必需的。我們可以將 Tensor.cpp 中的代碼包裝在一個文件中，然后為每個類型重復使用。我們將代碼分割成單獨文件的原因是這樣可以加快編譯速度。第二，當我們談論類型替換(例如用浮點數代替 real)時，我們的意思是，C 預處理器將在編譯期執行這些替換。并且在預處理之前這些嵌入源代碼的宏定義都沒有什么弊端。
通用構建(第二部分)
我們現在有所有的 Tensor 類型的源文件，我們需要考慮如何創建相應的頭文件聲明，以及如何將 THTensor_(方法)和 THPTensor_(方法)轉化成 TH
- THP_API PyObject * THPTensor_(New)(THTensor *ptr);
我們使用相同的策略在頭文件的源文件中生成代碼。在 csrc/Tensor.h 中，我們執行以下操作：
- #include "generic/Tensor.h"
- #include <TH/THGenerateAllTypes.h>
- #include "generic/Tensor.h"
- #include <TH/THGenerateHalfType.h>
從通用的頭文件中抽取代碼和用相同的宏定義包裝每個類型具有同樣的效果。唯一的區別就是前者編譯后的代碼包含在同一個頭文件中，而不是分為多個源文件。
***，我們需要考慮如何「轉換」或「替代」函數類型。如果我們查看相同的頭文件，我們會看到一堆 #define 語句，其中包括：
- #define THPTensor_(NAME) TH_CONCAT_4(THP,Real,Tensor_,NAME)
這個宏表示，源代碼中的任何匹配形如 THPTensor_(NAME)的字符串都應該替換為 THPRealTensor_NAME，其中 Real 參數是從符號 Real 所在的 #define 定義的時候派生的。因為我們的頭文件代碼和源代碼都包含所有上述類型的宏定義，所以在預處理器運行之后，生成的代碼就是我們想要的。
TH 庫中的代碼為 THTensor_(NAME)定義了相同的宏，支持這些功能的轉移。如此一來，我們最終就會得到帶有專用代碼的頭文件和源文件。
#### 模塊對象和類型方法，我們現在已經看到如何在 THP 中封裝 TH 的 Tensor 定義，并生成了 THPFloatTensor_init(...)等 THP 方法。現在我們可以從我們創建的模塊中了解上面的代碼實際上做了什么。THPTensor_(init)中的關鍵行是：
- # THPTensorBaseStr, THPTensorType are also macros that are specific
- # to each type
- PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);
該函數將 Tensor 對象注冊到擴展模塊，因此我們可以在我們的 Python 代碼中使用 THPFloatTensor，THPIntTensor 等。
只是單純的創建 Tensors 不是很有用 - 我們需要能夠調用 TH 定義的所有方法。以下是一個在 Tensor 上調用就地(in-place)zero_ 方法的簡單例子。
- x = torch.FloatTensor(10)
- x.zero_()
我們先看看如何向新定義的類型中添加方法。「類型對象」中的有一個字段 tp_methods。此字段包含方法定義數組(PyMethodDefs)，用于將方法(及其底層 C / C ++實現)與類型相關聯。假設我們想在我們的 PyFloatObject 上定義一個替換該值的新方法。我們可以按照下面的步驟來實現這一想法：
- static PyObject * replace(PyFloatObject *self, PyObject *args) {
- double val;
- if (!PyArg_ParseTuple(args, "d", &val))
- return NULL;
- self->ob_fval = val;
- Py_RETURN_NONE
- }
Python 版本的等價方法
- def replace(self, val):
- self.ob_fval = fal
閱讀更多的關于在 CPython 中如何定義方法頗具啟發性。通常，方法將對象的實例作為***個參數，以及可選的位置參數和關鍵字參數。這個靜態函數是在我們的浮點數上注冊為一個方法：
- static PyMethodDef float_methods[] = {
- {"replace", (PyCFunction)replace, METH_VARARGS,
- "replace the value in the float"
- },
- {NULL} /* Sentinel */
- }
這會注冊一個名為 replace 的方法，該方法由同名的 C 函數實現。METH_VARARGS 標志表示該方法使用包含函數所有參數的參數元組。該元組設置為類型對象的 tp_methods 字段，然后我們可以對該類型的對象使用 replace 方法。
我們希望能夠在 THP 張量等價類上調用所有的 TH 張量的方法。然而，為所有 TH 方法編寫封裝性價比極低。我們需要一個更好的方式來滿足這一需求。
PyTorch cwrap
PyTorch 實現自己的 cwrap 工具來包裝用于 Python 后端的 TH Tensor 方法。我們使用自定義 YAML 格式(http://yaml.org (http://yaml.org/))來定義包含一系列 C 方法聲明的.cwrapfile 文件。cwrap 工具獲取此文件，并以與 THPTensor Python 對象和 Python C 擴展方法調用相兼容的格式輸出包含打包方法的.cpp 源文件。此工具不僅用于生成包含 TH 的代碼，還包含 CuDNN。它是一款設計為可擴展的工具。
用于就地 addmv_功能的示例 YAML「聲明」如下：
- [[
- name: addmv_
- cname: addmv
- return: self
- arguments:
- - THTensor* self
- - arg: real beta
- default: AS_REAL(1)
- - THTensor* self
- - arg: real alpha
- default: AS_REAL(1)
- - THTensor* mat
- - THTensor* vec
- ]]
cwrap 工具的架構非常簡單。它先讀入一個文件，然后使用一系列插件進行處理。
源代碼在一系列的編譯通過時生成。首先，YAML「聲明」被解析和處理。然后，通過參數檢查和提取后源代碼逐個生成，定義方法頭，調用底層庫(如 TH)。***，cwrap 工具允許一次處理整個文件。addmv_的結果輸出可以在這里找到：
https://gist.github.com/killeent/c00de46c2a896335a52552604cc4d74b.
為了與 CPython 后端進行交互，該工具生成一個 PyMethodDefs 數組，可以存儲或附加到 THPTensor 的 tp_methods 字段。
在包裝 Tensor 方法的具體情況下，構建過程首先從 TensorMethods.cwrap 生成輸出源文件。該源文件就是通用 Tensor 源文件中的 #include 后面的文件。所有這些都發生在預處理器執行之前。結果，所有生成的方法包裝器都執行與上述 THPTensor 代碼相同的運作過程。因此，單個通用聲明和定義也適用于其它類型。
合而為一
到目前為止，我們已經展示了如何擴展 Python 解釋器來創建一個新的擴展模塊，如何定義我們新的 THPTensor 類型，以及如何為所有與 TH 連接的類型的 Tensor 生成源代碼。簡單來說，我們將染指匯編。
Setuptool 允許我們定義一個用于編譯的擴展模塊。整個 torch._C 擴展模塊文件是通過收集所有源文件、頭文件、庫等，并創建一個 setuptool 擴展來編譯的。然后，由 setuptool 處理構建擴展模塊本身。我將在隨后的一篇博文中探討更多的構建過程。
總而言之，讓我們回顧一下我們的四個問題：
1. PyTorch 如何通過擴展 Python 解釋器來定義可以從 Python 代碼中調用的 Tensor 類型?
它使用 CPython 的框架來擴展 Python 解釋器并定義新的類型，同時尤其關注為所有類型生成代碼。
2. PyTorch 如何封裝實際定義 Tensor 屬性和方法的 C 的類庫?
它通過定義一個由 TH Tensor 支持的新型 THPTensor。再通過 CPython 后端的各種語法規則，函數調用信息就會轉發到這個張量。
3. PyTorch 的 C 類包裝器如何生成 Tensor 方法的代碼?
它需要我們提供自定義的 YAML 格式的代碼，并通過使用多個插件通過一系列處理步驟來為每個方法生成源代碼。
4. PyTorch 的編譯系統如何編譯這些組件并生成可運行的應用程序?
它需要一堆源/頭文件、庫和編譯指令來構建使用 Setuptool 的擴展模塊。
本博文只是 PyTorch 構建系統的部分概述。還有更多的細節，但我希望這是對 Tensor 類的多數組件的通用介紹。
資源：
https://docs.python.org/3.7/extending/index.html 對于理解如何編寫 Python 的 C / C++擴展模塊***價值。
原文：https://gist.github.com/killeent/4675635b40b61a45cac2f95a285ce3c0
【本文是51CTO專欄機構“機器之心”的原創譯文，微信公眾號“機器之心( id: almosthuman2014)”】
