原标题:教程 | PyTorch内部机制解析:如何通过PyTorch实现Tensor
选自Github
机器之心编译
参与:朱乾树、黄小天
PyTorch 中的基本单位是张量(Tensor)。本文的主旨是如何在 PyTorch 中实现 Tensor 的概述,以便用户可从 Python shell 与之交互。本文主要回答以下四个主要问题:
1. PyTorch 如何通过扩展 Python 解释器来定义可以从 Python 代码中调用的 Tensor 类型?
2. PyTorch 如何封装实际定义 Tensor 属性和方法的 C 的类库?
3. PyTorch 的 C 类包装器如何生成 Tensor 方法的代码?
4. 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 函数来实现这些方法。该对象还可能包含表示其状态所需的任意字段。
定义新类型的准则如下:
1. 创建一个结构体,它定义了新对象将包括的属性
2. 定义类型的类型对象
结构体本身可能十分简单。在 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_deal