PyTorch学习总结(六)——Tensor实现

---

1. Python的C扩展

其实只要你懂得C语言编程，给Python添加新的内置(build-in)模块将十分容易。这些扩展(extension)模块可以实现两种无法直接在Python中进行的操作：他们可以实现新的内置对象类型，以及可以调用C语言的库函数和进行系统调用。

为了支持扩展，Python API定义了一个函数(functions)、宏命(macros)令和变量(variables)的集合，该集合提供了对Python运行时(run-time)系统的多方面的访问。Python API可以通过包含头文件Python.h的方式，整合进C语言源文件中。

注意： C语言的扩展接口指的是CPython，且扩展模块在其他Python实现上无效。在许多情况下，可以避免编写C扩展，并保留可移植性到其他实现上。例如，如果你的用例调用C库函数或进行系统调用，你可以考虑使用ctypes模块或cffi库而不是编写自定义的C代码。这些模块允许您编写Python代码来与C代码进行对接，并且Python实现比编写和编译C扩展模块更方便。

A Simple Example

我们准备构建一个名为spam的扩展模块，并假设我们想要为C库函数system()创建一个Python接口。该函数以null结尾(null-terminated)的字符串作为参数并返回一个整数。我们希望这个函数可以通过下列形式在Python中进行调用：

import spamstatus = spam.system("ls -l")

我们首先新建一个spammodule.c文件。(一般来说，如果一个模块的名字叫spam，则实现它的C语言文件应该命名为spammodule.c；如果模块的名字非常长，就像spammify，那对应的文件名就可以直接写为spammify.c`。)

该文件的第一行为：

#include <Python.h>

这可以获取Python API。

注意：因为Python会定义一些预处理器的定义，这些定义会影响某些系统的标准头文件，所以你必须在文件的一开始包含Python.h文件。

然后添加函数的主体代码：

static PyObject * spam_system(PyObject *self, PyObject *args)  {      const char *command;      int sts;      if (!PyArg_ParseTuple(args, "s", &command))          return NULL;      sts = system(command);      return PyLong_FromLong(sts);  }  

其中PyArg_ParseTuple(args, “s”, &command)检测参数列表是否存在错误，错误则返回NULL，否则把指令传给command。

现在相当于我们已经实现了spam_system函数，接下来我们就要讨论如何让Python能调用它。

在这里我们的模块名称为spam，而spam_system是它的方法。所以我们需要把spam_system添加进spam的方法列表中。

static PyMethodDef SpamMethods[] = {      ...      {"system",  spam_system, METH_VARARGS,       "Execute a shell command."},      ...      {NULL, NULL, 0, NULL}        /* Sentinel */  };  

其中，METH_VARARGS是参数传递的标准形式,它通过Python的元组在Python解释器和C函数之间传递参数。若采用METH_KEYWORD方式，则Python解释器和C函数之间将通过Python的字典类型在两者之间进行参数传递。system是函数spam_system在模块中的名字。

现在方法有了，我们需要构建一个表征该模块的结构体，然后将方法列表添加进去。如下所示：

static struct PyModuleDef spammodule = {     PyModuleDef_HEAD_INIT,     "spam",   /* name of module */     spam_doc, /* module documentation, may be NULL */     -1,       /* size of per-interpreter state of the module,                 or -1 if the module keeps state in global variables. */     SpamMethods  };  

定义完结构体后，就要定义初始化函数了。如下所示：

PyMODINIT_FUNC PyInit_spam(void)  {      return PyModule_Create(&spammodule);  }  

注意：在初始化函数中，一定要传入模块结构体。用PyModule_Create()函数初始化模块结构体。并且在Python 3中，初始化函数一定要用PyInit_name()的方式命名。初始化函数没有添加static声明，所以模块的初始化函数是模块的唯一对外接口。

上述材料都准备好后，我们接下来讨论如何将生成这个模块。我们编写一个setup.py文件，声明Extension，并用setuptools工具对模块进行打包。代码如下：

from setuptools import setup, Extension, distutils, Command, find_packagesC = Extension("spam",              libraries=main_libraries,              sources=['spammodule.c'],              language='c',              )setup(name="spam",       version=1.0,      ext_modules=[C],      )

打开终端，cd到setup.py目录下，接着输入下面两条指令即可完成模块的编译。然后就可以用import指令对spam模块进行导入啦~

python setup.py buildpython setup.py install 

在编译的过程中PyMODINIT_FUNC方法被调用，完成了spam的定义。

总结：让我们再来回顾一下整个过程：构建spammodule.c文件->#include <Python.h>及编写方法函数->构建某块结构体及初始化函数->构建setup.py文件，声明Extension及打包模块。

2. PyTorch中的C拓展

_C模块的构建

在PyTorch中，很多类都会继承_C模块的内容。_C模块中包含了Tensor、Storage等常用类型的C语言实现。Tensor类型可以看做是_C模块中的一种类型对象，它们是包含与被包含的关系。

首先，我们先来看_C模块的实现。这个过程与第一部分讲的很类似。_C模块的主体声明定义代码在torch/csrc/Module.cpp中。该文件的一行也是#include <Python.h>，函数的方法实现有很多，这里先不讨论。接下来我们找到模块结构体的定义及初始化函数的声明。如下所示：

#if PY_MAJOR_VERSION == 2PyMODINIT_FUNC init_C()#elsePyMODINIT_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));    ...}    

注意：不同版本的Python，初始化函数的命名方式不同。源码中用#if、endif方式进行区分。为了方便理解，以Python 3为例。代码整理如下：

PyMODINIT_FUNC PyInit__C(){    ...    static struct PyModuleDef torchmodule = {        PyModuleDef_HEAD_INIT,        "torch._C",        NULL,        -1,        methods.data()    };    ASSERT_TRUE(module = PyModule_Create(&torchmodule));    ...    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));    ...}    

我们可以看到这里的声明和定义方式与之前有所不同。结构体的定义放在了初始化函数里，模块名称为torch._C。接着我们打开setup.py文件。

...main_sources = [    "torch/csrc/PtrWrapper.cpp",    "torch/csrc/Module.cpp",    "torch/csrc/Generator.cpp",    "torch/csrc/Size.cpp",    "torch/csrc/Exceptions.cpp",    "torch/csrc/Storage.cpp",    "torch/csrc/DynamicTypes.cpp",    "torch/csrc/byte_order.cpp",    "torch/csrc/utils.cpp",    "torch/csrc/expand_utils.cpp",    "torch/csrc/utils/invalid_arguments.cpp",    ...]main_sources += split_types("torch/csrc/Tensor.cpp")...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'],      )

可以看到和上一节的过程基本一样，这里就不再赘述了。在编译的过程中。Module.cpp中的PyMODINIT_FUNC函数被调用，完成torch._C的定义，及各种类型Tensor的初始化函数的调用。如：ASSERT_TRUE(THPDoubleTensor_init(module))。

注意：上述代码中的main_sources里还包含了Tensor.cpp这个文件。它主要包含了多类Tensor的实现。

Tensor的实现

仔细阅读Module.cpp和Tensor.cpp文件，我们发现好像并没有DoubleTensor的初始化方法。这就是源码的厉害之处了。源码利用宏定义的方式，实现Tensor的多态。接下来就让我们一步步解开它的神秘面纱吧~

首先，我们注意到setup.py中有这么一句话main_sources += split_types("torch/csrc/Tensor.cpp")。这里调用了tools/setup_helpers/split_types.py文件中的方法，该方法主要实现了两个步骤：

1. 重新命名Tensor.cpp为Tensor<Type>.cpp格式，并保存在torch\csrc\generated\中；

2. 将Tensor.cpp的最后一行：

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

改为

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

得到如下文件：

这里比较重要的是#include "TH/THGenerate<Type>Type.h"，我们打开其中一个文件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

其中real和Real这两个宏定义很重要，在后面通过宏定义实现Tensor多态中会广泛用到。

现在我们通过main_sources += split_types("torch/csrc/Tensor.cpp")已经实现了多类Tensor.cpp的文件的生成。但是这些文件里面好像并没有什么实际的代码。其实具体的代码实现在这一句话中：#define TH_GENERIC_FILE "torch/src/generic/Tensor.cpp"。generic/Tensor.cpp里面几乎实现了Tensor的所有操作。比如：THPTensor_(init)等。在generic/Tensor.cpp中我们会看到大量THPTensor_(name)形式的函数，这些函数名其实是宏定义。打开torch\csrc\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_*也是宏定义，实现字符串的拼接功能。我们以上面的THPTensor_(init)函数为例，其类型为THPTensor_(NAME)，经过拼接后得到的函数名为THPRealTensor_init。记得之前提到过的宏定义Real吗？它有Int、Float、Long等对应类型，每一类型就可以实现对应的函数:THPIntTensor_init、THPFloatTensor_init等。这不就是之前

ASSERT_TRUE(THPLongTensor_init(module));ASSERT_TRUE(THPIntTensor_init(module));ASSERT_TRUE(THPShortTensor_init(module));

中的函数吗？他们都来源于同一个函数THPTensor_(init)，这就是PyTorch中多态的实现。

现在我们再来看看函数THPTensor_(init)：

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

THPTensorType.tp_methods = THPTensor_(methods);相当于给THPTensorType对象添加方法。

这里THPTensorBaseStr是一个宏定义：#define THPTensorBaseStr TH_CONCAT_STRING_2(Real,TensorBase)。而PyModule_AddObject(module, THPTensorBaseStr, (PyObject *)&THPTensorType);就相当于给module添加了一个新的对象THPTensorType。

那么THPTensorType是什么东西呢？其实它是THPTensor对象的一个类型对象，同时也是一个结构体#define THPTensorType TH_CONCAT_3(THP,Real,TensorType)。我们可以在generic\Tensor.h中找到THPTensor的定义：

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

一个PyObject_HEAD头，一个THTensor类型指针指向具体内容。根据宏定义THPTensor可以转换成THPIntTensor、THPLongTensor等类型。但他们要想实现被调用，就得给他们声明一个对象。其中包含对象名，方法等内容。

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_getattr */  0,                                     /* tp_setattr */  0,                                     /* tp_reserved */  0,                                     /* tp_repr */  0,                                     /* tp_as_number */  0,                                     /* tp_as_sequence */  &THPTensor_(mappingmethods),           /* 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 | Py_TPFLAGS_BASETYPE,   /* tp_flags */  NULL,                                  /* tp_doc */  0,                                     /* tp_traverse */  0,                                     /* tp_clear */  0,                                     /* tp_richcompare */  0,                                     /* tp_weaklistoffset */  0,                                     /* tp_iter */  0,                                     /* tp_iternext */  0,   /* will be assigned in init */    /* tp_methods */  0,   /* will be assigned in init */    /* tp_members */  0,                                     /* tp_getset */  0,                                     /* tp_base */  0,                                     /* tp_dict */  0,                                     /* tp_descr_get */  0,                                     /* tp_descr_set */  0,                                     /* tp_dictoffset */  0,                                     /* tp_init */  0,                                     /* tp_alloc */  THPTensor_(pynew),                     /* tp_new */};

注意：这里的格式与第一节中模块的结构体很相似。不要混淆了，这里是声明的类型对象，不是模块的结构体。

以THPIntTensor类型为例，经过PyModule_AddObject函数，将其对象IntTensorBase添加进torch._C模块中。后面我们就可以用_C.IntTensorBase的方式访问THPIntTensor类型了。