pytorch源碼：C拓展

Python C拓展

和普通的CPython方式拓展類似，下面主要看pytorch中的拓展子產品“_C"的定義和相應其他子產品的添加方式。pytorch中的拓展子產品定義代碼主要在torch/csrc/Module.cpp中，在預備知識中熟悉Python如何拓展C後，直接在Module.cpp找到感興趣的初始化部分：

#if PY_MAJOR_VERSION == 2
PyMODINIT_FUNC init_C()
#else
PyMODINIT_FUNC PyInit__C()
#endif
{
...
#if PY_MAJOR_VERSION == 2
  ASSERT_TRUE(module = Py_InitModule("torch._C", methods.data()));
#else
  static struct PyModuleDef torchmodule = {
     PyModuleDef_HEAD_INIT,
     "torch._C",
     NULL,
     -1,
     methods.data()
  };
  ASSERT_TRUE(module = PyModule_Create(&torchmodule));
#endif
  ...

  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));
...
}

在編譯過程中PyMODINIT_FUNC方法被調用，完成了"torch._C"的定義，接着就是各種類型Tensor的初始化函數調用，該部分在後面詳細來看。

和普通C拓展套路一緻，最終在編譯階段的setup.py檔案中,聲明Extension 執行setup加入拓展和用到的lib:

C = Extension("torch._C",
              libraries=main_libraries,
              sources=main_sources,
              language='c++',
              extra_compile_args=main_compile_args + extra_compile_args,
              include_dirs=include_dirs,
              library_dirs=library_dirs,
              extra_link_args=extra_link_args + main_link_args + [make_relative_rpath('lib')],
              )
extensions.append(C)
...
...
setup(name="torch", version=version,
      description="Tensors and Dynamic neural networks in Python with strong GPU acceleration",
      ext_modules=extensions,
      cmdclass=cmdclass,
      packages=packages,
      package_data={'torch': [
          'lib/*.so*', 'lib/*.dylib*',
          'lib/torch_shm_manager',
          'lib/*.h',
          'lib/include/TH/*.h', 'lib/include/TH/generic/*.h',
          'lib/include/THC/*.h', 'lib/include/THC/generic/*.h',
          'lib/include/ATen/*.h',
      ]},
      install_requires=['pyyaml', 'numpy'],
      )

代碼生成

一開始閱讀代碼，并沒有發現IntTensor部分的實作，如上面PyMODINIT_FUNC函數中的THPIntTensor_init(module)，和最開始提到的IntTensorBase。另一個奇怪的是很多檔案同時在根目錄（pytorch/torch/csrc/）和generic目錄下出現。以根目錄下Tensor.cpp為例：

#include <Python.h>
#include <structmember.h>

#define THP_HOST_HALF

#include <stdbool.h>
#include <vector>
#include <stack>
#include <tuple>
#include <TH/THMath.h>

#include "torch/csrc/THP.h"
#include "torch/csrc/copy_utils.h"
#include "torch/csrc/DynamicTypes.h"

//generic_include TH torch/csrc/generic/Tensor.cpp

這個cpp檔案和一般的cpp檔案不同，注意最後一句，在pytorch的setup.py中準備source清單時會調用：

main_sources += split_types("torch/csrc/Tensor.cpp")

其中split_types在pytorch/tools/setup_helpers/目錄下實作，主要實作了兩個功能：

重命名Tensor.cpp為Tensor[Type].cpp，其中Type為Float、Int等
各個Tensor[Type].cpp内容的最後一行内容改變為：

#define TH_GENERIC_FILE "torch/src/generic/Tensor.cpp"
#include "TH/THGenerate[Type]Type.h"

在lib/TH下則可以看到對應的THGenerate[Type]Type.h，如TH/THGenerateIntType.h：

#ifndef TH_GENERIC_FILE
#error "You must define TH_GENERIC_FILE before including THGenerateIntType.h"
#endif

#define real int32_t
#define ureal uint32_t
#define accreal int64_t
#define TH_CONVERT_REAL_TO_ACCREAL(_val) (accreal)(_val)
#define TH_CONVERT_ACCREAL_TO_REAL(_val) (real)(_val)
#define Real Int
#define THInf INT_MAX
#define TH_REAL_IS_INT
#line 1 TH_GENERIC_FILE
#include TH_GENERIC_FILE
#undef real
#undef ureal
#undef accreal
#undef Real
#undef THInf
#undef TH_REAL_IS_INT
#undef TH_CONVERT_REAL_TO_ACCREAL
#undef TH_CONVERT_ACCREAL_TO_REAL

#ifndef THGenerateManyTypes
#undef TH_GENERIC_FILE
#endif

頭檔案中有兩個重要的宏定義：

#define real int32_t
#define Real Int

而在對應的Tensor.h中則有下面的宏定義:

#ifndef THP_TENSOR_INC
#define THP_TENSOR_INC

#define THPTensor                   TH_CONCAT_3(THP,Real,Tensor)
#define THPTensorStr                TH_CONCAT_STRING_3(torch.,Real,Tensor)
#define THPTensorClass              TH_CONCAT_3(THP,Real,TensorClass)
#define THPTensor_(NAME)            TH_CONCAT_4(THP,Real,Tensor_,NAME)
......

其中TH_CONCAT_*同樣是宏定義，實作拼接字元串的功能，如對于IntTensor類型會有：

THPTensor_(init) will be convert to THPIntTensor_init

而類似THPTensor_(NAME)格式的字串在generic目錄下的代碼實作中大量出現。

這裡的命名值得一提，源碼中會經常遇到THP和TH字首的變量，前者是pytorch中的變量字首，差別于後者，源自Torch庫中的對應變量。pytorch底層實作中調用了大量的Torch庫。

THPTensor實作

在弄清楚上面兩部分之後，再看generic目錄下的代碼就清晰很多了，還是以THPIntTensor為例來看。這裡的THPIntTensor實際上是pytorch拓展的一個新Python類型。如果接觸過Python源碼的話會很清楚，定義一個新類型需要：

定義該對象包括哪些東西
為該對象定義類型

下面通過對比Python和pytorch的對象機制來簡單說明：

Python對象機制

以C實作的Python為例，對于int類型，需要為其定義該類型：

typedef struct tagPyIntObject
{
    PyObject_HEAD;
    int value;
} PyIntObject;

對應類型有：

PyTypeObject PyInt_Type =
{
     PyObject_HEAD_INIT(&PyType_Type),
     "int",
     ...
};

其中PyObject_HEAD為宏定義，定義了所有對象所共有的部分，包括對象的引用計數和對象類型等共有資訊，這也是Python中多态的來源。PyObject_HEAD_INIT是類型初始化的宏定義，簡單來看如下：

#define PyObject_HEAD \
 int refCount;\
 struct tagPyTypeObject *type

 #define PyObject_HEAD_INIT(typePtr)\
 0, typePtr

如果對Python源碼很感興趣，強烈推薦陳儒（Robert Chen）的《Python源碼剖析》，很是精煉。按照書中的介紹，我這裡也實作了一個傻瓜Python以供參考：https://github.com/zqhZY/smallpy

pytorch對象機制

pytorch拓展的Tensor類型與Python的一般類型的定義類似，generic目錄下的Tensor.h中有如下定義：

struct THPTensor {
  PyObject_HEAD
  // Invariant: After __new__ (not __init__), this field is always non-NULL.
  THTensor *cdata;
};

同樣的簡潔明了，一個PyObject_HEAD頭，一個THTensor類型指針指向具體内容。這裡的THTensor類型就涉及到了TH庫，也是源碼的最底層的C語言實作。TH庫的内容到最後再看，這裡簡單了解為一個類型。同樣值得一提，根據上面的内容，這裡的THPTensor和THTensor最終會轉換成不同的類型：如Int對應THPIntTensor和THIntTensor。

同樣，對于THPTensor對象也對應一個類型對象的定義，generic/Tensor.cpp中有：

PyTypeObject THPTensorType = {
  PyVarObject_HEAD_INIT(NULL, 0)
  "torch._C." THPTensorBaseStr,          /* tp_name */
  sizeof(THPTensor),                     /* tp_basicsize */
  0,                                     /* tp_itemsize */
  (destructor)THPTensor_(dealloc),       /* tp_dealloc */
  0,                                     /* tp_print */
  ...
  ...
  0,                                     /* tp_alloc */
  THPTensor_(pynew),                     /* tp_new */
};

結構體中包括了很多指針，這裡看最後的THPTensor_(pynew)，該方法在該類型對象建立時調用，對應Python層面的____new____函數。找到該函數：

static PyObject * THPTensor_(pynew)(PyTypeObject *type, PyObject *args, PyObject *kwargs)
{
    ...
    // torch.Tensor()
    if (num_args == 0) {
     self->cdata = THPTensor_(_new)();
     return (PyObject*)self.release();
    }
    ...
}

THPTensor_(pynew)先是為THPTensor申請記憶體，之後根據參數不同建立并初始化。

這裡簡單看代碼中不傳參數情況下建立Tensor調用THPTensor_(_new):

static THTensor* THPTensor_(_new)()
{
  THTensorPtr tensor(THTensor_(new)(LIBRARY_STATE_NOARGS));
  if (!tensor->storage) {
    tensor->storage = THStorage_(new)(LIBRARY_STATE_NOARGS);
  }
  return tensor.release();
}

該函數傳回THTensor類型，方法中調用的THTensor_(new)和THStorage_(new)均為TH庫中的内容，實作底層封裝和記憶體申請。其中Storage儲存了Tensor的值，之後具體看。

Python C拓展Tensor類型

看完 Tensor類型的初始化後，接下來是看如何把各種Tensor類型加入到”_C“子產品下供上層Python調用。這裡回到一開始的torch/csrc/Module.cpp中PyMODINIT_FUNC方法的一系列初始化：

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));

該部分初始化對應到generic/Tensor.cpp中的THPTensor_(init):

bool THPTensor_(init)(PyObject *module)
{
  ...
  THPTensorType.tp_methods = THPTensor_(methods);
  ...
  PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);
  THPTensor_(initCopyMethods)();
  return true;
}

該段代碼有兩點需要解釋，一是Tensor子產品的添加，二是Tensor對象的方法集的指定。

Tensor子產品添加

PyModule_AddObject函數為Python C拓展API：

int PyModule_AddObject(PyObject *module, const char *name, PyObject *value)
Add an object to module as name.

而THPTensorBaseStr則是另外一個宏定義：

#define THPTensorBaseStr            TH_CONCAT_STRING_2(Real,TensorBase)

同樣是字元串拼接宏,對不同類型THPTensorBaseStr最終轉換成[Type]TensorBase，再回到文章最開始的：

class IntTensor(_C.IntTensorBase, _TensorBase)

由此得到了Python層可以繼承的_C.IntTensorBase。

Tensor類型添加方法

參考Pyhton實作，在定義一個對象後，對應的類型結構體中包含一個指針，指向該類型可以調用的方法集，例如list類型的用法：

a = []
a.append(1)

在pytorch的Tensor類型中該指針即為tp_methods，該指針的指派，如上面THPTensor_(init)中:

THPTensorType.tp_methods = THPTensor_(methods);

奇怪的是THPTensor_(methods)并不能在源碼中找到，而在generic/methods可以看到一些.cwrap檔案，如對ones函數有：

[[
  name: ones
  variants:
    - function
  auto_gpu: False
  return: argument 0
  arguments:
    - arg: THTensor* result
      output: True
    - arg: THSize* size
      long_args: True
]]

pytorch中自己實作了插件，根據這種YAML格式文本對TH庫代碼進行封裝生成對應代碼，插件代碼在tools/cwrap/plugins。如addmv_函數對應生成的内容：

https://gist.github.com/killeent/c00de46c2a896335a52552604cc4d74b 。

小結

到此，主要說明了pytorch中Tensor類型的定義及其子產品拓展機制，可以使上層的Python調用C拓展的類型和相應方法。可以看到，pytorch中使用了代碼生成方式，隻定義一個模闆，不同類型的Tensor對象通過該模闆生成，避免了大量重複代碼，雖然一開始一頭霧水，但确實比較巧妙。

篇幅原因，這裡并沒有深入去看TH庫部分的代碼，pytorch對torch庫做了CPython類封裝，重用了大量代碼，TH中主要的一個部分是THTensor的實作，後面再繼續整理TH部分的代碼。

參考文獻

Python拓展C

Python源碼剖析精簡版

A Tour of PyTorch Internals (Part I)

A quick tour of Torch internals

作者：zqh_zy

連結：https://www.jianshu.com/p/91af2ab867d5

來源：簡書

簡書著作權歸作者所有，任何形式的轉載都請聯系作者獲得授權并注明出處。

pytorch源碼：C拓展

Python C拓展

代碼生成

THPTensor實作

Python C拓展Tensor類型

小結

繼續閱讀

來自python的【條件控制/語句循環/break/continue/else/pass】一、條件控制二、語句循環

無法解析的外部符号 wmain，該符号在函數 "void cdecl mainCRTStartupHelper(struct HINSTANCE *,unsigned short con......

TestLink導出用例轉換工具(XML2Excel)

YAML簡介和PyYAML安全操作YAML支援的類型YAML的優點：yaml的基本文法python操作

Small tricks

libsvm for python 安裝

學習軟體測試基礎測試第七天

Zeppelin 配置通路 REST APIApache Zeppelin Configuration REST API

【Torch】最簡潔logging使用指南

27. Remove Element(清單)題目代碼

Cloud Studio初體驗

使用 ctypes 進行 Python 和 C 的混合程式設計

【python】【資料處理】畫多元資料分布圖

【python】netconf協定對接管理裝置

「Python 網絡自動化」NETCONF —— Python 使用 NETCONF 管理配置 H3C 網絡裝置

在python中建立excel并寫入