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 spam
status = spam.system("ls -l")我们首先新建一个spammodule.c文件。(一般来说，如果一个模块的名字叫spam，则实现它的C语言文件应该命名为spammodule.c；如果模块的名字非常长，就像spammify，那对应的文件名就可以直接写为spammify.c`。)
该文件的第一行为：
#include 这可以获取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_packages
C = 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 build
python setup.py install 在编译的过程中PyMODINIT_FUNC方法被调用，完成了spam的定义。
总结：让我们再来回顾一下整个过程：构建
spammodule.c文件->#include->构建某块结构体及初始化函数->构建及编写方法函数 setup.py文件，声明Extension及打包模块。
2. PyTorch中的C拓展
_C模块的构建
在PyTorch中，很多类都会继承_C模块的内容。_C模块中包含了Tensor、Storage等常用类型的C语言实现。Tensor类型可以看做是_C模块中的一种类型对象，它们是包含与被包含的关系。
首先，我们先来看_C模块的实现。这个过程与第一部分讲的很类似。_C模块的主体声明定义代码在torch/csrc/Module.cpp中。该文件的一行也是#include ，函数的方法实现有很多，这里先不讨论。接下来我们找到模块结构体的定义及初始化函数的声明。如下所示：
#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));
...
} 注意：不同版本的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文件中的方法，该方法主要实现了两个步骤：
重新命名
Tensor.cpp为Tensor格式，并保存在.cpp torch\csrc\generated\中；将
Tensor.cpp的最后一行：
//generic_include TH torch/csrc/generic/Tensor.cpp改为
#define TH_GENERIC_FILE "torch/src/generic/Tensor.cpp"
#include "TH/THGenerateType.h" 得到如下文件：
这里比较重要的是#include "TH/THGenerate，我们打开其中一个文件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类型了。