PyTorch內(nèi)部機制解析：如何通過PyTorch實現(xiàn)Tensor
PyTorch 中的基本單位是張量(Tensor)。本文的主旨是如何在 PyTorch 中實現(xiàn) Tensor 的概述，以便用戶可從 Python shell 與之交互。本文主要回答以下四個主要問題：
- PyTorch 如何通過擴展 Python 解釋器來定義可以從 Python 代碼中調(diào)用的 Tensor 類型?
- PyTorch 如何封裝實際定義 Tensor 屬性和方法的 C 的類庫?
- PyTorch 的 C 類包裝器如何生成 Tensor 方法的代碼?
- PyTorch 的編譯系統(tǒng)如何編譯這些組件并生成可運行的應用程序?
擴展 Python 解釋器
PyTorch 定義了一個新的包 torch。本文中，我們將考慮._C 模塊。這是一個用 C 編寫的被稱為「擴展模塊」的 Python 模塊，它允許我們定義新的內(nèi)置對象類型(例如 Tensor)和調(diào)用 C / C ++函數(shù)。
._C 模塊定義在 torch/csrc/Module.cpp 文件中。init_C()/ PyInit__C()函數(shù)創(chuàng)建模塊并根據(jù)需要添加方法定義。這個模塊被傳遞給一些不同的__init()函數(shù)，這些函數(shù)會添加更多的對象到模塊中，以及注冊新的類型等。
__init() 可調(diào)用的部分函數(shù)如下：
- 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()函數(shù)將每種類型的 Tensor 對象添加到._C 模塊，以便它們可以在._C 模塊中調(diào)用。下面我們來了解這些方法的工作原理。
THPTensor 類型
PyTorch 很像底層的 TH 和 THC 類庫，它定義了一個專門針對多種不同的類型數(shù)據(jù)的「通用」Tensor。在考慮這種專業(yè)化的工作原理之前，我們首先考慮如何在 Python 中定義新的類型，以及如何創(chuàng)建通用的 THPTensor 類型。
Python 運行時將所有 Python 對象都視為 PyObject * 類型的變量，PyObject * 是所有 Python 對象的「基本類型」。每個 Python 類型包含對象的引用計數(shù)，以及指向?qū)ο蟮摹割愋蛯ο蟆沟闹羔?。類型對象確定類型的屬性。例如，該對象可能包含一系列與類型相關(guān)聯(lián)的方法，以及調(diào)用哪些 C 函數(shù)來實現(xiàn)這些方法。該對象還可能包含表示其狀態(tài)所需的任意字段。
定義新類型的準則如下：
- 創(chuàng)建一個結(jié)構(gòu)體，它定義了新對象將包括的屬性
- 定義類型的類型對象
結(jié)構(gòu)體本身可能十分簡單。在 Python 中，實際上所有浮點數(shù)類型都是堆上的對象。Python float 結(jié)構(gòu)體定義為：
- typedef struct {
- PyObject_HEAD
- double ob_fval;
- } PyFloatObject;
PyObject_HEAD 是引入實現(xiàn)對象的引用計數(shù)的代碼的宏，以及指向相應類型對象的指針。所以在這種情況下，要實現(xiàn)浮點數(shù)，所需的唯一其他「狀態(tài)」是浮點值本身。
現(xiàn)在，我們來看看 THPTensor 類型的結(jié)構(gòu)題：
- struct THPTensor {
- PyObject_HEAD
- THTensor *cdata;
- };
很簡單吧?我們只是通過存儲一個指針來包裝底層 TH 張量。關(guān)鍵部分是為新類型定義「類型對象」。我們的 Python 浮點數(shù)的類型對象的示例定義的形式如下：
- 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 字段設(shè)置為 sizeof(PyFloatObject)。這是為了讓 Python 知道 PyFloatObject 調(diào)用 PyObject_New()時需要分配多少內(nèi)存。你可以設(shè)置的字段的完整列表在 CPython 后端的 object.h 中定義：
https://github.com/python/cpython/blob/master/Include/object.h.
THPTensor 的類型對象是 THPTensorType，它定義在 csrc/generic/Tensor.cpp 文件中。該對象定義了 THPTensor 的類型名稱、大小及映射方法等。
我們來看看我們在 PyTypeObject 中設(shè)置的 tp_new 函數(shù)：
- PyTypeObject THPTensorType = {
- PyVarObject_HEAD_INIT(NULL, 0)
- ...
- THPTensor_(pynew), /* tp_new */
- };
tp_new 函數(shù)可以創(chuàng)建對象。它負責創(chuàng)建(而不是初始化)該類型的對象，相當于 Python 中的__new()__方法。C 實現(xiàn)是一個靜態(tài)方法，該方法傳遞實例化的類型和任意參數(shù)，并返回一個新創(chuàng)建的對象。
- 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
我們的新函數(shù)的***件事就是為 THPTensor 分配內(nèi)存。然后，它會根據(jù)傳遞給該函數(shù)的參數(shù)進行一系列的初始化。例如，當從另一個 THPTensor y 創(chuàng)建 THPTensor x 時，我們將新創(chuàng)建的 THPTensor 的 cdata 字段值設(shè)置為以 y 的底層 TH Tensor 作為參數(shù)并調(diào)用 THTensor_(newWithTensor)返回的結(jié)果。這一過程中有內(nèi)存大小、存儲、NumPy 數(shù)組和序列的類似的構(gòu)造函數(shù)。
注意，我們只使用了 tp_new 函數(shù)，而不是同時使用 tp_new 和 tp_init(對應于 Python 中的 __init()__函數(shù))。
Tensor.cpp 中定義的另一個重要的部分是索引的工作原理。PyTorch Tensors 支持 Python 的映射協(xié)議。這樣我們可以做如下事情：
- 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(標量)，或更新現(xiàn)有 Tensor 的值。閱讀這些實現(xiàn)代碼，以更好地了解 PyTorch 是如何支持基本張量索引的。
通用構(gòu)建(***部分)
我們可以花費大量時間探索 THPTensor 的各個方面，以及它如何與一個新定義 Python 對象相關(guān)聯(lián)。但是我們?nèi)匀恍枰靼? THPTensor_(init)()函數(shù)是如何轉(zhuǎn)換成我們在模塊初始化中使用的 THPIntTensor_init()函數(shù)。我們該如何使用定義「通用」Tensor 的 Tensor.cpp 文件，并使用它來生成所有類型序列的 Python 對象?換句話說，Tensor.cpp 里遍布著如下代碼：
- return THPTensor_(New)(THTensor_(new)(LIBRARY_STATE_NOARGS));
這說明了我們需要使類型特定的兩種情況：
- 我們的輸出代碼將調(diào)用 THP
Tensor_New(...)代替調(diào)用 THPTensor_(New) - 我們的輸出代碼將調(diào)用 TH
Tensor_new(...)代替調(diào)用 THTensor_(new)
換句話說，對于所有支持的 Tensor 類型，我們需要「生成」已經(jīng)完成上述替換的源代碼。這是 PyTorch 的「構(gòu)建」過程的一部分。PyTorch 依賴于配置工具(https://setuptools.readthedocs.io/en/latest/)來構(gòu)建軟件包，我們在頂層目錄中定義一個 setup.py 文件來自定義構(gòu)建過程。
使用配置工具構(gòu)建擴展模塊的一個組件是列出編譯中涉及的源文件。但是，我們的 csrc/generic/Tensor.cpp 文件未列出!那么這個文件中的代碼最終是如何成為最終產(chǎn)品的一部分呢?
回想前文所述，我們從以上的 generic 目錄中調(diào)用 THPTensor *函數(shù)(如 init)。如果我們來看一下這個目錄，會發(fā)現(xiàn)一個定義了的 Tensor.cpp 文件。此文件的***一行很重要：
- //generic_include TH torch/csrc/generic/Tensor.cpp
請注意，雖然這個 Tensor.cpp 文件被 setup.py 文件引用，但它被包裝在一個叫 Python helper 的名為 split_types 的函數(shù)里。這個函數(shù)需要輸入一個文件，并在該文件內(nèi)容中尋找「//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 定義為一個浮點數(shù)，所以泛型 Tensor 實現(xiàn)中的任何代碼將指向一個 real 對象，實際上 real 被替換為浮點數(shù)。在對應的文件 THGenerateIntType.h 中，同樣的宏定義將用 int 替換 real。
這些輸出文件從 split_types 返回，并添加到源文件列表中，因此我們可以看到不同的類型的.cpp 代碼是如何創(chuàng)建的。
這里需要注意以下幾點：***，split_types 函數(shù)不是必需的。我們可以將 Tensor.cpp 中的代碼包裝在一個文件中，然后為每個類型重復使用。我們將代碼分割成單獨文件的原因是這樣可以加快編譯速度。第二，當我們談論類型替換(例如用浮點數(shù)代替 real)時，我們的意思是，C 預處理器將在編譯期執(zhí)行這些替換。并且在預處理之前這些嵌入源代碼的宏定義都沒有什么弊端。
通用構(gòu)建(第二部分)
我們現(xiàn)在有所有的 Tensor 類型的源文件，我們需要考慮如何創(chuàng)建相應的頭文件聲明，以及如何將 THTensor_(方法)和 THPTensor_(方法)轉(zhuǎn)化成 TH
- THP_API PyObject * THPTensor_(New)(THTensor *ptr);
我們使用相同的策略在頭文件的源文件中生成代碼。在 csrc/Tensor.h 中，我們執(zhí)行以下操作：
- #include "generic/Tensor.h"
- #include <TH/THGenerateAllTypes.h>
- #include "generic/Tensor.h"
- #include <TH/THGenerateHalfType.h>
從通用的頭文件中抽取代碼和用相同的宏定義包裝每個類型具有同樣的效果。唯一的區(qū)別就是前者編譯后的代碼包含在同一個頭文件中，而不是分為多個源文件。
***，我們需要考慮如何「轉(zhuǎn)換」或「替代」函數(shù)類型。如果我們查看相同的頭文件，我們會看到一堆 #define 語句，其中包括：
- #define THPTensor_(NAME) TH_CONCAT_4(THP,Real,Tensor_,NAME)
這個宏表示，源代碼中的任何匹配形如 THPTensor_(NAME)的字符串都應該替換為 THPRealTensor_NAME，其中 Real 參數(shù)是從符號 Real 所在的 #define 定義的時候派生的。因為我們的頭文件代碼和源代碼都包含所有上述類型的宏定義，所以在預處理器運行之后，生成的代碼就是我們想要的。
TH 庫中的代碼為 THTensor_(NAME)定義了相同的宏，支持這些功能的轉(zhuǎn)移。如此一來，我們最終就會得到帶有專用代碼的頭文件和源文件。
#### 模塊對象和類型方法，我們現(xiàn)在已經(jīng)看到如何在 THP 中封裝 TH 的 Tensor 定義，并生成了 THPFloatTensor_init(...)等 THP 方法?，F(xiàn)在我們可以從我們創(chuàng)建的模塊中了解上面的代碼實際上做了什么。THPTensor_(init)中的關(guān)鍵行是：
- # THPTensorBaseStr, THPTensorType are also macros that are specific
- # to each type
- PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);
該函數(shù)將 Tensor 對象注冊到擴展模塊，因此我們可以在我們的 Python 代碼中使用 THPFloatTensor，THPIntTensor 等。
只是單純的創(chuàng)建 Tensors 不是很有用 - 我們需要能夠調(diào)用 TH 定義的所有方法。以下是一個在 Tensor 上調(diào)用就地(in-place)zero_ 方法的簡單例子。
- x = torch.FloatTensor(10)
- x.zero_()
我們先看看如何向新定義的類型中添加方法?！割愋蛯ο蟆怪械挠幸粋€字段 tp_methods。此字段包含方法定義數(shù)組(PyMethodDefs)，用于將方法(及其底層 C / C ++實現(xiàn))與類型相關(guān)聯(lián)。假設(shè)我們想在我們的 PyFloatObject 上定義一個替換該值的新方法。我們可以按照下面的步驟來實現(xiàn)這一想法：
- 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
閱讀更多的關(guān)于在 CPython 中如何定義方法頗具啟發(fā)性。通常，方法將對象的實例作為***個參數(shù)，以及可選的位置參數(shù)和關(guān)鍵字參數(shù)。這個靜態(tài)函數(shù)是在我們的浮點數(shù)上注冊為一個方法：
- static PyMethodDef float_methods[] = {
- {"replace", (PyCFunction)replace, METH_VARARGS,
- "replace the value in the float"
- },
- {NULL} /* Sentinel */
- }
這會注冊一個名為 replace 的方法，該方法由同名的 C 函數(shù)實現(xiàn)。METH_VARARGS 標志表示該方法使用包含函數(shù)所有參數(shù)的參數(shù)元組。該元組設(shè)置為類型對象的 tp_methods 字段，然后我們可以對該類型的對象使用 replace 方法。
我們希望能夠在 THP 張量等價類上調(diào)用所有的 TH 張量的方法。然而，為所有 TH 方法編寫封裝性價比極低。我們需要一個更好的方式來滿足這一需求。
PyTorch cwrap
PyTorch 實現(xiàn)自己的 cwrap 工具來包裝用于 Python 后端的 TH Tensor 方法。我們使用自定義 YAML 格式(http://yaml.org (http://yaml.org/))來定義包含一系列 C 方法聲明的.cwrapfile 文件。cwrap 工具獲取此文件，并以與 THPTensor Python 對象和 Python C 擴展方法調(diào)用相兼容的格式輸出包含打包方法的.cpp 源文件。此工具不僅用于生成包含 TH 的代碼，還包含 CuDNN。它是一款設(shè)計為可擴展的工具。
用于就地 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 工具的架構(gòu)非常簡單。它先讀入一個文件，然后使用一系列插件進行處理。
源代碼在一系列的編譯通過時生成。首先，YAML「聲明」被解析和處理。然后，通過參數(shù)檢查和提取后源代碼逐個生成，定義方法頭，調(diào)用底層庫(如 TH)。***，cwrap 工具允許一次處理整個文件。addmv_的結(jié)果輸出可以在這里找到：
https://gist.github.com/killeent/c00de46c2a896335a52552604cc4d74b.
為了與 CPython 后端進行交互，該工具生成一個 PyMethodDefs 數(shù)組，可以存儲或附加到 THPTensor 的 tp_methods 字段。
在包裝 Tensor 方法的具體情況下，構(gòu)建過程首先從 TensorMethods.cwrap 生成輸出源文件。該源文件就是通用 Tensor 源文件中的 #include 后面的文件。所有這些都發(fā)生在預處理器執(zhí)行之前。結(jié)果，所有生成的方法包裝器都執(zhí)行與上述 THPTensor 代碼相同的運作過程。因此，單個通用聲明和定義也適用于其它類型。
合而為一
到目前為止，我們已經(jīng)展示了如何擴展 Python 解釋器來創(chuàng)建一個新的擴展模塊，如何定義我們新的 THPTensor 類型，以及如何為所有與 TH 連接的類型的 Tensor 生成源代碼。簡單來說，我們將染指匯編。
Setuptool 允許我們定義一個用于編譯的擴展模塊。整個 torch._C 擴展模塊文件是通過收集所有源文件、頭文件、庫等，并創(chuàng)建一個 setuptool 擴展來編譯的。然后，由 setuptool 處理構(gòu)建擴展模塊本身。我將在隨后的一篇博文中探討更多的構(gòu)建過程。
總而言之，讓我們回顧一下我們的四個問題：
1. PyTorch 如何通過擴展 Python 解釋器來定義可以從 Python 代碼中調(diào)用的 Tensor 類型?
它使用 CPython 的框架來擴展 Python 解釋器并定義新的類型，同時尤其關(guān)注為所有類型生成代碼。
2. PyTorch 如何封裝實際定義 Tensor 屬性和方法的 C 的類庫?
它通過定義一個由 TH Tensor 支持的新型 THPTensor。再通過 CPython 后端的各種語法規(guī)則，函數(shù)調(diào)用信息就會轉(zhuǎn)發(fā)到這個張量。
3. PyTorch 的 C 類包裝器如何生成 Tensor 方法的代碼?
它需要我們提供自定義的 YAML 格式的代碼，并通過使用多個插件通過一系列處理步驟來為每個方法生成源代碼。
4. PyTorch 的編譯系統(tǒng)如何編譯這些組件并生成可運行的應用程序?
它需要一堆源/頭文件、庫和編譯指令來構(gòu)建使用 Setuptool 的擴展模塊。
本博文只是 PyTorch 構(gòu)建系統(tǒng)的部分概述。還有更多的細節(jié)，但我希望這是對 Tensor 類的多數(shù)組件的通用介紹。
資源：
https://docs.python.org/3.7/extending/index.html 對于理解如何編寫 Python 的 C / C++擴展模塊***價值。
原文：https://gist.github.com/killeent/4675635b40b61a45cac2f95a285ce3c0
【本文是51CTO專欄機構(gòu)“機器之心”的原創(chuàng)譯文，微信公眾號“機器之心( id: almosthuman2014)”】
