选自blog.christianperone，作者：Christian S. Perone，机器之心编译，原文链接：
blog.christianperone.com/2018/03/pyt…
作为 Facebook 人工智能团队（FAIR）提供支持的深度学习框架，PyTorch 自 2017 年 1 月推出以来立即成为了一种流行开发工具。其在调试、编译等方面的优势使其受到了学界研究者们的普遍欢迎。本文中，来自蒙特利尔综合理工学院的研究员 Christian S. Perone 将为我们介绍这种神经网络框架的内部架构，揭开 PyTorch 方便好用的真正原因。
前言
本文主要介绍了 PyTorch 代码库，旨在为 PyTorch 及其内部架构设计提供指导，核心目标是为那些想了解 API 知识之外的人提供有益的帮助，并给出之前教程所没有的新内容。注意：PyTorch 构建系统需要大量使用代码设置，因此其他人的描述我将不再重复。如果你感兴趣，请参考原文提供的扩展资料。
C/C++中 Python 扩展对象的简介
你可能知道可以借助 C/C++扩展 Python，并开发所谓的「扩展」。PyTorch 的所有繁重工作由 C/C++实现，而不是纯 Python。为了定义 C/C++中一个新的 Python 对象类型，你需要定义如下实例的一个类似结构：
// Python object that backs torch.autograd.Variable
struct THPVariable {
PyObject_HEAD
torch::autograd::Variable cdata;
PyObject* backward_hooks;
};
如上，在定义的开始有一个称之为 PyObject_HEAD 的宏，其目标是标准化 Python 对象，并扩展至另一个结构，该结构包含一个指向类型对象的指针，以及一个带有引用计数的字段。
Python API 中有两个额外的宏，分别称为 Py_INCREF() 和 Py_DECREF()，可用于增加和减少 Python 对象的引用计数。多实体可以借用或拥有其他对象的引用（因此引用计数被增加），而只有当引用计数达到零，Python 才会自动删除那个对象的内存。想了解更多有关 Python C/++扩展的知识，请参见：https://docs.python.org/3/extending/newtypes.html。
有趣的事实：使用小的整数作为索引、计数等在很多应用中非常见。为了提高效率，官方 CPython 解释器缓存从-5 到 256 的整数。正由于此，声明 a = 200; b = 200; a is b 为真，而声明 a = 300; b = 300; a is b 为假。
Zero-copy PyTorch 张量到 Numpy，反之亦然
PyTorch 有专属的张量表征，分离 PyTorch 的内部表征和外部表征。但是，由于 Numpy 数组的使用非常普遍，尤其是当数据加载源不同时，我们确实需要在 Numpy 和 PyTorch 张量之间做转换。正由于此，PyTorch 给出了两个方法（from_numpy() 和 numpy()），从而把 Numpy 数组转化为 PyTorch 数组，反之亦然。如果我们查看把 Numpy 数组转化为 PyTorch 张量的调用代码，就可以获得有关 PyTorch 内部表征的更多洞见：
at::Tensor tensor_from_numpy(PyObject* obj) {
if (!PyArray_Check(obj)) {
throw TypeError("expected np.ndarray (got %s)", Py_TYPE(obj)->tp_name);
}
auto array = (PyArrayObject*)obj;
int ndim = PyArray_NDIM(array);
auto sizes = to_aten_shape(ndim, PyArray_DIMS(array));
auto strides = to_aten_shape(ndim, PyArray_STRIDES(array));
// NumPy strides use bytes. Torch strides use element counts.
auto element_size_in_bytes = PyArray_ITEMSIZE(array);
for (auto& stride : strides) {
stride /= element_size_in_bytes;
}
// (...) - omitted for brevity
void* data_ptr = PyArray_DATA(array);
auto& type = CPU(dtype_to_aten(PyArray_TYPE(array)));
Py_INCREF(obj);
return type.tensorFromBlob(data_ptr, sizes, strides, [obj](void* data) {
AutoGIL gil;
Py_DECREF(obj);
});
}
代码摘自（tensor_numpy.cpp：https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_numpy.cpp#L88）
正如你在这段代码中看到的，PyTorch 从 Numpy 表征中获取所有信息（数组元数据），并创建自己的张量。但是，正如你从被标注的第 18 行所看到的，PyTorch 保留一个指向内部 Numpy 数组原始数据的指针，而不是复制它。这意味着 PyTorch 将拥有这一数据，并与 Numpy 数组对象共享同一内存区域。
还有一点很重要：当 Numpy 数组对象越出范围并获得零引用（zero reference）计数，它将被当作垃圾回收并销毁，这就是为什么 Numpy 数组对象的引用计数在第 20 行有增加。该行之后，PyTorch 将从这一 Numpy 数据 blob 中创建一个新的张量对象，并且在创建这一新张量的过程中，PyTorch 将会传递内存数据指针，连同内存大小、步幅以及稍后张量存储将会使用的函数（我们将会在下节讨论），从而通过减少 Numpy 数组对象的引用计数并使 Python 关心这一对象内存管理而释放数据。
tensorFromBlob() 方法将创建一个新张量，但只有在为这一张量创建一个新「存储」之后。存储是指存储数据指针的地方，它并不在张量结构内部。张量存储正是我们下一节要讨论的内容。
张量存储
张量的实际原始数据并不是立即保存在张量结构中，而是保存在我们称之为「存储（Storage）」的地方，它是张量结构的一部分。
正如我们前面在 tensor_from_numpy() 中看到的代码，它调用了 tensorFromBlob() 函数以从原始数据 Blob 中创建一个张量。tensorFromBlob() 函数在内部会调用另一个名为 storageFromBlob() 函数，该函数主要根据类型为数据创建一个存储。例如在 CPU 浮点型的情况下，它会返回一个新的 CPUFloatStorage 实例。
CPUFloatStorage 基本上是包含 utility 函数的包装类（wrapper），且实际存储结构如下所示称为 THFloatStorage：
typedef struct THStorage
{
real *data;
ptrdiff_t size;
int refcount;
char flag;
THAllocator *allocator;
void *allocatorContext;
struct THStorage *view;
} THStorage;
如上所示，THStorage 有一个指向原始数据、原始数据大小、flags 和 allocator 的指针，我们会在后面详细地讨论它们。值得注意的是，THStorage 不包含如何解释内部数据的元数据，这是因为存储对保存的内容「无处理信息的能力」，只有张量才知道如何「查看」数据。
因此，你可能已经意识到多个张量可以指向相同的存储，而仅仅对数据采用不同的解析。这也就是为什么我们以不同的形状或维度，查看相同元素数量的张量会有很高的效率。下面的 Python 代码表明，在改变张量的形状后，存储中的数据指针将得到共享。
>>> tensor_a = torch.ones((3, 3))
>>> tensor_b = tensor_a.view(9)
>>> tensor_a.storage().data_ptr() == tensor_b.storage().data_ptr()
True
如 THFloatStorage 结构中的第七行代码所示，它有一个指向 THAllocator 结构的指针。它因为给分配器（allocator）带来灵活性而显得十分重要，其中 allocator 可以用来分配存储数据。
typedef struct THAllocator
{
void* (*malloc)(void*, ptrdiff_t);
void* (*realloc)(void*, void*, ptrdiff_t);
void (*free)(void*, void*);
} THAllocator;
代码摘自（THAllocator.h：https://github.com/pytorch/pytorch/blob/master/aten/src/TH/THAllocator.h#L16）
如上所述，该结构有三个函数指针字段来定义分配器的意义：malloc、realloc 和 free。对于分配给 CPU 的内存，这些函数当然与传统的 malloc/realloc/free POSIX 函数相关。然而当我们希望分配存储给 GPU，我们最终会使用如 cudaMallocHost() 那样的 CUDA 分配器，我们可以在下面的 THCudaHostAllocator malloc 函数中看到这一点。
static void *THCudaHostAllocator_malloc(void* ctx, ptrdiff_t size) {
void* ptr;
if (size < 0) THError("Invalid memory size: %ld", size);
if (size == 0) return NULL;
THCudaCheck(cudaMallocHost(&ptr, size));
return ptr;
}
代码摘自（THCAllocator.c：https://github.com/pytorch/pytorch/blob/master/aten/src/THC/THCAllocator.c#L3）
如上所示，分配器调用了一个 cudaMallocHost() 函数。你可能已经注意到版本库组织中有缩写的表示模式，在浏览版本库时记住这些约定非常重要，它们在 PyTorch README 文件中有所总结：
-
TH = TorcH
-
THC = TorcH Cuda
-
THCS = TorcH Cuda Sparse
-
THCUNN = TorcH CUda Neural Network
-
THD = TorcH Distributed
-
THNN = TorcH Neural Network
-
THS = TorcH Sparse
该约定同样存在于函数/类别名和其它对象中，因此了解它们十分重要。你可以在 TH 代码中找到 CPU 分配器，在 THC 代码中找到 CUDA 分配器。最后，我们可以看到主张量 THTensor 结构的组成：
typedef struct THTensor
{
int64_t *size;
int64_t *stride;
int nDimension;
THStorage *storage;
ptrdiff_t storageOffset;
int refcount;
char flag;
} THTensor;
如上，THTensor 的主要结构为张量数据保留了 size/strides/dimensions/offsets/等，同时还有存储 THStorage。我们可以将所有这些结构总结为以下图表：
现在，如果我们有多重处理的需求，且希望在多个不同的进程中共享张量数据，那么我们需要一个共享内存的方法。否则每次另一个进程需要张量或我们希望实现 Hogwild 训练过程以将所有不同的进程写入相同的内存区域时，我们就需要在进程间创建副本，这是非常低效的。因此，我们将在下一节讨论共享内存的特定存储方法。
共享内存
共享内存可以用很多种不同的方法实现（依赖于支持的平台）。PyTorch 支持部分方法，但为了简单起见，我将讨论在 MacOS 上使用 CPU（而不是 GPU）的情况。由于 PyTorch 支持多种共享内存的方法，由于代码中包含很多级的间接性，这部分会有点困难。
PyTorch 为 Python multiprocessing 模块提供了一个封装器，可以从 torch.multiprocessing 导入。他们对该封装器中的实现做出了一些变动，以确保每当一个 Tensor 被放在队列上或和其它进程共享时，PyTorch 可以确保仅有一个句柄的共享内存会被共享，而不会共享 Tensor 的完整新副本。现在，很多人都不知道 PyTorch 中的 Tensor 方法是 share_memory_()，然而，该函数正好可以触发那个特定 Tensor 的保存内存的完整重建。该方法的执行过程是创建共享内存的一个区域，其可以在不同的进程中使用。最终，该函数可以调用以下的函数：
static THStorage* THPStorage_(newFilenameStorage)(ptrdiff_t size)
{
int flags = TH_ALLOCATOR_MAPPED_SHAREDMEM | TH_ALLOCATOR_MAPPED_EXCLUSIVE;
std::string handle = THPStorage_(__newHandle)();
auto ctx = libshm_context_new(NULL, handle.c_str(), flags);
return THStorage_(newWithAllocator)(size, &THManagedSharedAllocator, (void*)ctx);
}
如上所示，该函数使用了一个特殊的分类器 THManagedSharedAllocator 来创建另一个存储。它首先定义了一些 flags，然后创建了一个格式为 /torch_ [process id] _ [random number] 的字符串句柄，最后在使用特殊的 THManagedSharedAllocator 创建新的存储。该分配器有一个指向 PyTorch 内部库 libshm 的函数指针，它将实现名为 Unix Domain Socket 的通信以共享特定 quyu 的内存句柄。这种分配器实际上是「smart allocator」的特例，因为它包含通信控制逻辑单元，并使用了另一个称之为 THRefcountedMapAllocator 的分配器，它将创建市级共享内存区域并调用 mmp() 以将该区域映射到进程虚拟地址空间。
