Pytorch底层源码解读（一）概览

前言

作为最受欢迎的深度学习框架，Pytorch如今已拥有极大的用户群体以及开发者。但对于开发者而言，针对日益臃肿的pytorch框架进一步更新迭代已经成为了较大的问题，特别是对刚想要上手对pytorch底层框架进行理解的初学者而言。
因此本系列主要针对于pytorch底层框架中的核心部分进行解读，为读者展现其背后工作机理的同时也能使得读者在阅读完本系列的文章后，能够对pytorch框架有个基本的了解，甚至可以做到对其进行修改、优化。
本文以pytorch 2.0.1作为参考进行讲解。

Pytorch框架概览

pytorch工作原理

绝大部分朋友应该都听说过前端、后端的概念，但是很多人对其的理解可能仅限于网页前端、服务器后端这类。实际上，前后端的概念可以进一步拓展。
在编译过程中，像llvm这类的工具可以将不同的语言如C++、Rust等转化成中间代码，这部分中间代码是与语言无关的，也就是说通过这类工具可以解决不同语言的异构问题，实现这个过程的工具在编译原理中称为前端。在我们拿到与语言无关的中间代码后，我们再通过一个工具将中间代码翻译为针对于不同指令集的汇编语言，这样便可以实现不同语言在不同cpu架构上运行了，而实现这个过程的工具在编译原理中就称为后端。

llvm.png

这里我们引用一下网上对于llvm的相关图片，可以从中看出llvm将不同的语言通过一系列处理生成了可以在不同的架构上执行的汇编语言，这里的前端就是C、C++等，后端就是llvm在不同架构上的backend。
事实上，pytorch框架也是借鉴了这个思路，针对于该框架而言，python便是前端，而用c++实现的框架完成了后端的工作，后面涉及到相关概念时再做具体说明。

目录树概览

打开pytorch源码，从目录树上可以看到，pytorch框架十分庞大，但实际上其中最核心的部分占比较小，对于入门而言，只需要掌握以下两个部分即可：

1. torch目录：这个部分主要存放的是是pytorch的前端代码，也包括了C++的实现，为了实现C++与python的混合开发，在这里使用了pybind、ctypes等python包。
2. aten目录：A Tensor Library的缩写，这个目录下主要包括了与Tensor相关的内容，例如Tensor的定义、存储、操作等。可以看到在aten/src/Aten目录下，算子实现都在native/目录中。其中有CPU的算子实现，以及CUDA的算子实现（cuda/）等。
   除了上述两个部分以外的其他目录，基本上都暂时不用了解，当然不是说这些目录不重要，比如c10、caffe2这些移植了caffe后端，但是我们在入门阶段不用去了解这些内容，把目光主要放在torch和aten两个目录即可。

torch解读

当然，细心的读者应该可以发现，torch目录下的文件结构有些眼熟，我们不妨打开安装好的pytorch包的目录，就会发现其实源代码下的torch目录下的文件与最终安装的pytorch包目录下几乎是一一对应的，除了少部分源代码中出现的.h头文件不存在以外。至于为什么会出现这个情况，写过python包的读者大概也了解，其实这就是python包（package）的结构。
要编写一个python包首先需要定义python包的目录结构，一个python包结构往往具有一个主模块和若干个子模块，然后通过__init__.py文件对某个目录进行进行标记，被标记的目录便识别为python包。我们也注意到，在torch文件夹下便存在这么一个文件，事实上，torch文件夹下的这部分代码就是pytorch的前端，用户在import torch的时候就是引入的这个目录。
对于相当一部分python包而言，它们的__init__.py文件为空，这也可以理解，因为该文件的主要作用是标记一个目录为python包，因此即使什么也不写仍然是可以的。但是也存在部分python包会在该文件下插入一部分代码，pytorch便是这么做的。
在讨论这么做的好处之前，我们先来引入一个新的概念：import关键字的执行原理。
首先我们要知道，python和c最大的不同在于python是解释型语言，而c是编译型语言（虽然官方的python底层是用c实现的），那么python中的import自然也不能与c中的#include划上等号。#include实际上只是很单纯地在编译期间将一个文件引入到了另一个文件里，而python的import却是一个运行时的概念。当python程序执行到某个文件的import时，如果需要导入的文件是一个.py文件，那么解释器会将其作为一个Module 对象进行导入。如果需要导入的是一个python包，即package，那么解释器会直接去执行需要导入的模块的__init__.py文件，然后将其作为一个 Module 对象给放在当前的全局变量中(通过globals()函数返回)，这样就实现了Module和package的统一。
后续的具体细节本文不再深究，我们先来看看torch\__init__.py里有什么。

__init__.py解读

首先是__all__部分，从中我们可以看到很多熟悉的字眼，如rand、is_tensor等，这些都是我们平时常用的函数。__all__ 是针对模块公开接口的一种约定，比起双下划线的方式（私有变量或者私有函数）， __all__以提供了”白名单“的形式暴露接口。具体的意思就是我们可以直接通过torch.xxx的方式，使用__all__中给定的函数。

__all__ = [
    'typename', 'is_tensor', 'is_storage', 'set_default_tensor_type',
    'set_rng_state', 'get_rng_state', 'manual_seed', 'initial_seed', 'seed',
    'save', 'load', 'set_printoptions', 'chunk', 'split', 'stack', 'matmul',
    'no_grad', 'enable_grad', 'rand', 'randn', 'inference_mode',
    'DoubleStorage', 'FloatStorage', 'LongStorage', 'IntStorage',
    'ShortStorage', 'CharStorage', 'ByteStorage', 'BoolStorage',
    '_TypedStorage',
    'DoubleTensor', 'FloatTensor', 'LongTensor', 'IntTensor',
    'ShortTensor', 'CharTensor', 'ByteTensor', 'BoolTensor', 'Tensor',
    'lobpcg', 'use_deterministic_algorithms',
    'are_deterministic_algorithms_enabled',
    'is_deterministic_algorithms_warn_only_enabled',
    'set_deterministic_debug_mode', 'get_deterministic_debug_mode',
    'set_float32_matmul_precision', 'get_float32_matmul_precision',
    'set_warn_always', 'is_warn_always_enabled',
]

加载库文件

接下来是第58行至第142行，这部分代码主要是处理在windows平台上运行pytorch会遇到的问题。在windows平台上，如果使用者创建了一个虚拟环境用于执行pytorch代码，那么在执行前需要将pytorch依赖的库目录插入到windows的dll搜索目录中，否则将会出现依赖问题。
一个题外话，代码中的win32其实指的并不是windows系统，而是其底层架构，不过不管是windows的64位系统还是32位系统，其底层架构均为win32，以免读者混淆，在此额外说明。
另外，关于这部分代码，对于没有涉及过python与C、C++混合编程的读者而言其实并不是很容易看懂，因为它使用了ctypes库，这是一个可以在python中调用由C、C++编写并导出的dll动态链接库的包。如下代码中给出的ctypes.CDLL('vcruntime140.dll')这部分，其实就是加载了用C、C++编写一个名为vcruntime140.dll的文件。另外，动态链接库在windows系统中的后缀是.dll，而在linux中的后缀是.so，pytorch代码中也加载了不少.so文件以适配linux系统，这在后面的解析中可以看到。关于ctypes的介绍网上有很多相关的资料，本文不再进行阐述。

if sys.platform == 'win32':
    pfiles_path = os.getenv('ProgramFiles', 'C:\\Program Files')
    ...
    try:
        ctypes.CDLL('vcruntime140.dll')
        ctypes.CDLL('msvcp140.dll')
        ctypes.CDLL('vcruntime140_1.dll')
    except OSError:
        print('''Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.
                 It can be downloaded at https://aka.ms/vs/16/release/vc_redist.x64.exe''')
    ...
    kernel32.SetErrorMode(prev_error_mode)

第146行到154行提供了根据不同的平台加载了一个叫做libtorch_global_deps的库的函数，这个函数将在下面进行调用。

def _load_global_deps():
    if platform.system() == 'Windows' or sys.executable == 'torch_deploy':
        return

    lib_name = 'libtorch_global_deps' + ('.dylib' if platform.system() == 'Darwin' else '.so')
    here = os.path.abspath(__file__)
    lib_path = os.path.join(os.path.dirname(here), 'lib', lib_name)

    ctypes.CDLL(lib_path, mode=ctypes.RTLD_GLOBAL)

第157行到202行的代码为加载libtorch_global_deps和其他一些库提供了一些方法。

if (USE_RTLD_GLOBAL_WITH_LIBTORCH or os.getenv('TORCH_USE_RTLD_GLOBAL')) and \
        platform.system() != 'Windows':
        ...
    from torch._C import *  # noqa: F403

细心的朋友之前或许已经发现了，在pytorch的源码中存在的torch\csrc并没有在安装好的torch目录下出现，这是因为这部分c++内容以编译好的动态库文件代替了，第204行到第207行就是将其引入的代码。

# Appease the type checker; ordinarily this binding is inserted by the
# torch._C module initialization code in C
if TYPE_CHECKING:
    import torch._C as _C

当然，实际上从55行到255行的代码对于我们理解pytorch的底层原理而言并不重要，一般而言这部分代码都不会进行改动，这部分代码我们可以统一的概括为根据不同平台加载pytorch所需要的适配的库文件。

定义基本工具

从261行到640行的代码实际上就是在__all__中给出的部分的实现，这部分代码被视为pytorch的一些基本工具。

def typename(o):
    if isinstance(o, torch.Tensor):
        return o.type()
        ...
    return _C._get_warnAlways()

不过这部分代码其实有相当一部分很难让对python了解不深入的人看懂，比如从379行到502行的use_deterministic_algorithms函数，这个函数用了几乎所有的篇幅写了一段很长的注释，但是具体的实现确实只用了一行代码。

def use_deterministic_algorithms(mode, *, warn_only=False):
    r""" Sets whether PyTorch operations must use "deterministic"
    algorithms. That is, algorithms which, given the same input, and when
    run on the same software and hardware, always produce the same output.
    When enabled, operations will use deterministic algorithms when available,
    and if only nondeterministic algorithms are available they will throw a
    :class:`RuntimeError` when called.
    ...
        # Backward mode nondeterministic error
        >>> torch.randn(10, requires_grad=True, device='cuda').index_select(0, torch.tensor([0], device='cuda')).backward()
        ...
        RuntimeError: index_add_cuda_ does not have a deterministic implementation...
    """
    _C._set_deterministic_algorithms(mode, warn_only=warn_only)

根据最后一行_C._set_deterministic_algorithms(mode, warn_only=warn_only)，可以很容易的知道这部分代码其实是调用了torch._C下的_set_deterministic_algorithms函数，但是当我们打开这个torch._C这个包对应的__init__.pyi文件时却发现这个函数是这样的：

def _set_deterministic_algorithms(mode: _bool, *, warn_only: _bool=...) -> None: ...  # THPModule_setDeterministicAlgorithms

好像什么也没有写。不过如果你往上翻代码，你会发现在669行有如下的一段注释：

# Defined in torch/csrc/Module.cpp
def _initExtension(shm_manager_path: str) -> None: ...  # THPModule_initExtension
...

毫无疑问，在这个.pyi文件中出现的函数，它们的具体实现其实是在torch/csrc文件夹下的，这也是为什么本文从一开始便提出仅需掌握这个目录下的文件即可。另外，.pyi文件是python中的类型提示文件，也被叫做存根文件stub file，用于提供代码的静态类型信息，也可以用来表示公共的接口。事实上这也很好理解，python是一个动态类型语言，它的类型推断是根据赋值操作实现的，而C、C++都是静态类型语言，它们需要给定变量或者函数的类型，因此.pyi文件给出了其静态类型，这样便做到了python和C、C++的绑定。关于其原理这里不再过多阐述，笔者将在这个系列后续的文章中做进一步解读和说明，现在我们只需要认为这部分代码通过C、C++实现了可以被python调用的函数即可。

定义数字常量

毫无疑问，这部分以拓展__all__的方式实现了对几个数字常量的定义。

from math import e , nan , inf , pi
__all__.extend(['e', 'pi', 'nan', 'inf'])

定义Storage和Tensor 类

从655行开始到760行的这部分代码就是重中之重了，它们的具体实现将在下一章进行讲解，本章仅针对于它们在__init__.py文件中起到的作用进行阐述。

from ._tensor import Tensor
from .storage import _StorageBase, _TypedStorage, _LegacyStorage, _UntypedStorage

# NOTE: New Storage classes should never be added. When adding a new
# dtype, use torch.storage._TypedStorage directly.

class ByteStorage(_LegacyStorage):
    @classproperty
    def dtype(self):
        return torch.uint8
...

_storage_classes = {
    _UntypedStorage, DoubleStorage, FloatStorage, LongStorage, IntStorage,
    ShortStorage, CharStorage, ByteStorage, HalfStorage, BoolStorage,
    QUInt8Storage, QInt8Storage, QInt32Storage, BFloat16Storage,
    ComplexFloatStorage, ComplexDoubleStorage, QUInt4x2Storage, QUInt2x4Storage,
    _TypedStorage
}

# The _tensor_classes set is initialized by the call to _C._initialize_tensor_type_bindings()
_tensor_classes: Set[Type] = set()

从上方给出的代码中可以看到，在本文件中引入了_tensor模块的Tensor类，又从storage模块引入了一些storage类，并实现了一些特定的storage类，诸如ByteStorage。其实从名称上可以看出该类是一个用于存储字节类型数据的存储对象，那么什么是存储对象呢？
实际上在pytorch中，每一个Tensor对象都对应着一个Storage对象，而每一个Storage对象都对应着一个或多个Tensor对象，这个知识点在pytorch的官方文档中其实是提到过的。Tensor对象好理解，因为所有使用者都使用过torch.tensor()或者torch.Tensor()这种函数创建过Tensor对象。tensor的中文名叫张量，这实际上是一个数学上的概念。正如本文在一开始提到的，Tensor相关的定义与实现都在aten目录下，在本章节中，我们不刻意去解读Tensor的具体的c++实现，我们将在下一节对其展开讲述。
我们根据导入的Tensor去查看这部分代码在python中的实现，从中我们可以看到，在_tensor.py这个文件中的第84行到1190行都是Tensor在python中的定义。

class Tensor(torch._C._TensorBase):
    def __deepcopy__(self, memo):
        if has_torch_function_unary(self):
            return handle_torch_function(Tensor.__deepcopy__, (self,), self, memo)
        if not self.is_leaf:
            raise RuntimeError("Only Tensors created explicitly by the user "
                               "(graph leaves) support the deepcopy protocol at the moment")
        ...
            device_type = DLDeviceType.kDLCPU
        else:
            raise ValueError('Unknown device type {} for Dlpack'.format(self.device.type))
        return (device_type, idx)

    __module__ = 'torch'

我们不细究Tensor当中的具体实现细节，先从宏观上来看，这个类首先继承了torch._C._TensorBase这个类，我们刚才提到torch._C是由C++代码生成的。因此可以知道的是，在python中的这个Tensor继承的是C++后端的TensorBase类的实现，而python中的这个Tensor类可以视为对C++中的TesorBase的再度封装，因为它在这里重写了一些python对象自带的函数，诸如深拷贝函数__deepcoy__。
值得注意的是，这些重写是很有必要的，比如我们在上文提到的一个Storage对应一个或多个Tensor，那么自然地，想要实现这个特性，我们就必须得重写__deepcopy__这个深拷贝函数，因此我们可以看到在第110行有这么一段代码：

       new_storage = self.storage().__deepcopy__(memo)

我们再来看Tensor类的storage方法：

    def storage(self):
        r"""
        storage() -> torch.Storage

        Returns the underlying storage.
        """
        if has_torch_function_unary(self):
            return handle_torch_function(Tensor.storage, (self,), self)

        return torch._TypedStorage(wrap_storage=self._storage(), dtype=self.dtype)

从上面两个代码中我们可以看到Tensor类的__deepcopy__方法其实是用的Storage对象的__deepcopy__方法，这也直接印证了一个Storage对象对应多个Tensor对象。然后我们再一层层的套娃去寻找，最后我们发现其实这个深拷贝函数返回的是这么一个对象：

class _TypedStorage:
    is_sparse = False

    dtype: torch.dtype

    def fill_(self, value):
        self[0:len(self)] = value
        return self
    ...

这个类并不是Storage的基类，但是它却有很多子类，接下来我们将来讲解这个类，至于Tensor类的剩下部分主要是对C++中的TensorBase的封装，用以实现在python端调用C++后端，这里就不再深入解读，毕竟真要详解，恐怕这篇文章的篇幅得翻几番，对于入门者想要了解pytorch框架而言，我们只需要知道python中的Tensor类实现了对后端的TensorBase的进一步封装即可。
首先我们看到storage.py这个文件，它是我们需要关注的重点，也是在__init__.py中引入过的部分（上文提到的）。我们看到storage.py的第16行到210行，这个部分是_StorageBase的定义，我们从名称可以看出者其实就是Storage的基类。

T = TypeVar('T', bound='Union[_StorageBase, _TypedStorage]')
class _StorageBase(object):
    _cdata: Any
    is_cuda: bool = False
    is_sparse: bool = False
    is_sparse_csr: bool = False
    device: torch.device

    def __init__(self, *args, **kwargs): ...  # noqa: E704
    ...

    def _untyped(self):
        return self

我们首先可以注意到的一点是，在这个类的最上方有一个泛型的类型注释，有的读者对TypeVar中的bound参数不了解，这其实意味着所有属于bound参数对应的类型及它的子类都可以通过类型检查，Union[_StorageBase, _TypedStorage]的意思就是既可以是_StorageBase类型，也可以是_TypedStorage类型。
在213行我们可以看到，它定义了一个新的类_UntypedStorage，从名称上我们可以知道其实就是和_TypedStorage相对的一个类，它继承了两个父类，但没有继续实现或定义一些成员方法，说明这两个父类提供的方法已经可以实现_UntypedStorage的要求了。

class _UntypedStorage(torch._C.StorageBase, _StorageBase):
    pass

我们再往下看，从第286行到第810行是_TypedStorage的定义，其实我们不难发现这个类的定义和_UntypedStorage还是有不小的出入的，它的定义相比于_UntypedStorage要更具象一些。

class _TypedStorage:
    is_sparse = False

    dtype: torch.dtype

    def fill_(self, value):
        self[0:len(self)] = value
        return self
        ...

从第824行开始到最后则是_LegacyStorage的定义，从中我们可以发现这个类实际上是继承了_TypedStorage。

class _LegacyStorage(_TypedStorage, metaclass=_LegacyStorageMeta):
    @classmethod
    def _new_shared(cls, size):
        """Creates a new storage in shared memory with the same data type"""
        untyped_storage = torch._UntypedStorage._new_shared(size * cls().element_size())
        return cls(wrap_storage=untyped_storage)

我们还记得我们是怎么一步步从__init__.py分析到这个阶段的嘛，现在让我们回到定义Storage和Tensor 类的这个部分，我们可以看到__init__.py导入的几个Storage类中，用的最多的其实就是_LegacyStorage了，下方是__init__.py的截取部分。

...
class QInt8Storage(_LegacyStorage):
    @classproperty
    def dtype(self):
        return torch.qint8

class QInt32Storage(_LegacyStorage):
    @classproperty
    def dtype(self):
        return torch.qint32

class QUInt4x2Storage(_LegacyStorage):
    @classproperty
    def dtype(self):
        return torch.quint4x2

class QUInt2x4Storage(_LegacyStorage):
    @classproperty
    def dtype(self):
        return torch.quint2x4

综上所述，我们可以得到如下的一个关系图，当然，这个继承图并不完整，比如对于Storage对象在不同的设备上的实现是不同的，因此后面还会有_CudaLegacyStorage这些类，它继承了_LegacyStorage类，实现了对GPU设备上的显存管理。

Storag继承图.png

现在我们可以对Tensor和Storage的实现思想有大致的了解了，实际上就是由一个python的基类继承C++的基类，然后再由这些基类去派生其他类，pytorch其他代码的思想也基本就是遵循这个理念。__init__.py后续的代码就是些细枝末节的东西了，这里不再赘述。

目录回顾与简介

通过对__init__.py的解读，我们已经对pytorch的框架的核心部分有了一定的了解，现在我们再来对pytorch源码的所有目录进行初步的了解（尽管当中很大一部分可能不是核心），但是本文将会对它们做一个简要的介绍，之后的文章将只会围绕当中的核心部分展开。

1.android是pytorch对安卓的支持库，准确地说是在安卓端编译pytorch需要用到的，写过安卓APP的应该都知道gradle这个构建工具，pytorch官方提供了支持gradle的pytorch编译方案。
2.aten文章最开始已经说过了，这里不赘述。
3.benchmarks该文件夹包含可生成各种 PyTorch 功能的可重复计时的脚本，它还提供了将 PyTorch 与其他框架进行比较的机制。
4.binaries这个目录其实挺让人迷惑的，看上去像是基准测试相关的工具文件。
5.c10，caffe2这两个目录下的文件在之前已经提到过了，它们移植了caffee的后端，关于caffee后端，本文不做讨论，读者可自行搜索。
6.cmake该目录下有许多.camke文件，本质上就是对pytorch的C++代码进行编译的文件。
7.docs该目录下存放了一些文档。
8.functorch，提到这个目录就得提到另一个项目JAX，本质上这个目录是pytorch框架对该项目一些功能的模仿。
9.ios作用类似于android目录。
10.moudles一些已经定义好了的模型。
11.mypy_plugins没啥用的目录。
12.scripts顾名思义，其中包含了很多sh脚本文件，作用面非常广泛，有在各种系统上进行编译的sh脚本，也有导出模型的脚本等等。
13.test存放了一堆测试文件，在对pytorch源码进行修改后，用于测试是否有效会用到。当然，大部分比较小的修改其实不需要使用这里的测试文件。
14.third-party第三方依赖。
15.utils一些常用的工具文件。
16.torch上文已经提到过。
17.torchgen用于生成代码。

总结

截止到这里，对pytorch框架的初步了解就到此结束，我们先是从__init__.py文件了解了pytorch框架的前端做了哪些核心的工作，然后对整个pytorch项目的目录进行了粗略的解释。下一篇文章将针对pytorch的Tensor展开讨论，主要解读其在后端所做的工作。