PyTorch源码解析（1）- 整体预览

> 第一篇文章节作为整体预览，一方面介绍代码的整体架构，一方面普及一些知识便于非ai开发人员也能快速了解相关背景。后续文章将针对不同的细节进行更深入的源码剖析。 > 本文已由公司账号公开发表在知乎：[https://zhuanlan.zhihu.com/p/598044604](https://zhuanlan.zhihu.com/p/598044604) [TOC] ## 1. 代码架构  本文截取的是`PyTorch 1.13`版本，PyTorch编译前的源码统计总计约213W行代码，编译后约300W行（编译时codegen会自动生成大量代码）。 其中我们选取了最重要的几个模块： 1. [`aten`]^(A TENsor library): A Tensor Library的缩写。与Tensor相关的内容都放在这个目录下。如Tensor的定义、存储、Tensor间的操作（即算子/OP）等； 可以看到在aten/src/Aten目录下，算子实现都在native/目录中。其中有CPU的算子实现，以及CUDA的算子实现（cuda/）等。 2. `torch`: 即PyTorch的前端代码。我们用户在import torch时实际引入的是这个目录。 其中包括前端的Python文件，也包括高性能的c++底层实现（csrc/）。为实现Python和c++模块的打通，这里使用了pybind作为胶水。在python中使用torch._C.[name]实际调用的就是libtorch.so中的c++实现，而PyTorch在前端将其进一步封装为python函数供用户调用。 3. [`c10`]^(Caffe TENsor library)、`caffe2`：移植caffe后端，c10指的是caffe tensor library，相当于caffe的aten。 PyTorch1.0完整移植了caffe2的源码，将两个项目进行了合并。引入caffe的原因是Pytorch本身拥有良好的前端，caffe2拥有良好的后端，二者在开发过程中拥有大量共享代码和库。例如在PyTorch中，c10::Device和at::Device是等价的。 4. `tools`：用于代码自动生成（codegen），例如autograd根据配置文件实现反向求导OP的映射。 5. `scripts`：一些脚本，用于不同平台项目构建或其他功能性脚本； __总结__：PyTorch源码架构中，最重要的两个目录是`aten`和`torch`目录。`aten`(A Tensor Library)目录主要是和`Tensor相关实现`的目录，包括`算子`的具体实现；而`torch`目录是`PyTorch前端`及其`底层实现`，用户`import torch`即安装的这个目录。这两个目录占据了除test/benchmark等绝大多数的源码行数。 ## 2. 前端&后端 ### 1. 概念 什么是PyTorch的前端、后端？ 在 PyTorch 中，`前端`指的是 PyTorch 的 `Python 接口`，它为用户提供了一组用于构建、训练和评估深度学习模型的工具。前端接口使用 Python 编写，因此可以轻松与其他 Python 库集成（如numpy等）。 `后端`指的是 PyTorch 的底层 `C++ 引擎`，它负责执行前端指定的计算。后端引擎使用张量表示计算图的节点和边，并使用高效的线性代数运算和卷积运算来执行计算。 后端引擎也负责与底层平台（如 GPU 和 CPU）进行交互，并将计算转换为底层平台能够执行的指令。这使得 PyTorch 能够在`多种平台`上运行，并能够利用 `GPU` 等加速设备加速计算。 ### 2. 前端 前端的使用即最普遍的PyTorch使用方法，在_pip install torch_后，可以在Python中`import torch`来使用： ```python #!/env/python import torch a = torch.rand(10) b = a + a print(b) ``` 我们可以在`site-packages`下查看torch包包含的内容：  对比下方`源码`的torch目录：  可以发现torch的wheel包安装的内容基本上就是torch目录下的python内容。而c++内容（csrc目录）并没有被复制过来，而是以编译好的动态库文件（\_C.cpython-*.so）代替。 在torch目录下的`__init__.py`文件中，可以看到将\_C(动态库)导入的地方。  ### 3. 后端 对于PyTorch的后端，也就是C++部分，实际上也是可以作为正常`c++ API`直接使用的。 除了可以直接使用_pip install torch_安装下的动态库文件+include目录来构建c++项目，也可以直接下载使用官方提供的`libtorch`。后者`不包含python`部分，因此可以在`嵌入式`设备等诸多场景下使用。 在使用PyTorch C++时，只需要将项目动态库与头文件正确链接即可。PyTorch官方提供了一个`libtorch项目编译`cmake[例子](https://pytorch.org/tutorials/advanced/cpp_export.html)： ```cmake cmake_minimum_required(VERSION 3.0 FATAL_ERROR) project(dcgan) find_package(Torch REQUIRED) add_executable(dcgan dcgan.cpp) target_link_libraries(dcgan "${TORCH_LIBRARIES}") set_property(TARGET dcgan PROPERTY CXX_STANDARD 14) ``` 这里的find\_package(Torch)这句是通过TorchConfig.cmake实现的库与头文件的链接，这是_PyTorch提供的cmake文件_，可以让我们快速实现c++项目编译。  torch的`c++ API`使用和python有一定区别，但在“形式”上大致相似，例如求tensor加法： ```cpp #include #include int main() { torch::Tensor a = torch::rand(10); auto b = a.add(a); std::cout << b << std::endl; return 0; } ``` 除此之外，Python中常用的API（如训练中的optimizer等）都可以通过c++ API实现；但是__部分只在Python中实现__的内容，例如module.to(device)是在python中递归执行.to()实现的，在c++ API中使用则会出现一定问题，因此使用时一定要注意。 ### 4. 前后端交互 接下来我们来看Python和C++黏合的部分（pybind）： 我们以Tensor为例。下面是PyTorch中的`Tensor实现`，我们可以看到它继承了\_C中的\_TensorBase。 {:class="width_60"} 在\_C目录下\_\_init\_\_.pyi中，我们可以看到\_TensorBase的类。.pyi文件是存根文件(stub file)，表示公共接口。其中有各种成员函数的定义，但是没有具体的实现。 在文件的注释中，我们可以看到其`具体实现`是在python_variable.cpp中。 {:class="width_60"} 在python\_variable.cpp文件中，我们可以看到PyTorch实际是用了`pybind`，将c++和Python进行交互的。 {:class="width_60"} 我们这里以Tensor.dtype为例，它的实现在THPVariable\_properties中。 找到这个结构体的实现，可以看到dtype具体绑定了THPVariable\_dtype()这个函数。因此在Python中执行tensor.dtype时，实际运行的就是这个c++函数。 {:class="width_60"} {:class="width_60"} __总结__：PyTorch`前端`主要是`python API`，在设计上采用`Pythonic`式的编程风格，可以让用户像使用python一样使用PyTorch；而`后端`主要指`C++ API`，其对外提供的C++接口，也可以一定程度上实现PyTorch的大部分功能，而且更适用于嵌入式等场景。而前端主要是通过`pybind`调用的后端c++实现，具体是c++被编译成`_C.[***].so动态库`，然后python调用`torch._C`实现调用c++中的函数。 ## 3. 动态图、静态图 一个主流的训练框架需要有`两大特征`： 1. 实现类似numpy的`张量计算`，可以使用`GPU`进行加速； 2. 实现带`自动微分系统`的深度神经网络。 而PyTorch能在一系列训练框架中脱颖而出成为今天的主流，主要原因是其原生支持`动态图`，具备对用户友好的特点。 ### 1. 概念 在`TensorFlow1.x`中，我们如果需要执行计算，需要建立一个`session`，并执行`session.run()`来执行。 其整个过程其实是将计算过程构成了一张`计算图`，然后运行这个图的根节点。这样先构成图，再运行图的方式我们称为`静态图`或者`图模式`。 {:class="width_60"} 而在PyTorch中，我们可以在计算的任意步骤`直接输出结果`。 原因是，PyTorch每一条语句是同步执行的，即每一条语句都是一个（或多个）`算子`，被`调用时实时执行`。这种实时执行算子的方式我们称为`动态图`或`算子模式`。 {:class="width_60"} __优缺点__ `动态图`的优点显而易见，可以兼容Python式的编程风格，实时打印编程结果，用户友好性做到最佳； `静态图`则在性能方面有一定优势，即在整个图执行前，可以将整张图进行编译优化，通过融合等策略改变图结构，从而实现较好的性能。 PyTorch原生支持动态图，但是也在静态图方面做了诸多尝试。我们将在下面一节简单介绍这些图模式。 ### 2. PyTorch图原理概述 本节简要讲解一下PyTorch动态图以及几种`静态图`实现的`原理`，这里主要是简要总结，不涉及源码。 后续将有其他的文章更详细地分析PyTorch各路径的源码及其实现。 PyTorch除了有原生的动态图之外，当前也在静态图方面有诸多路线尝试：分别有`torchscript`(`jit.script`、`jit.trace`)、`TorchDynamo`、`torch.fx`、`LazyTensor`。 这里用一幅图总结了各路线所处的层次：  首先PyTorch的动态图是从Python源码下降，拆分成多个python算子调用，具体调用到tensor的OP。经过pybind转换到c++，并通过dispatch机制选择不同设备下的算子实现，最终实现调用的底层设备实现（如Nvidia cudnn、intel mkl等）。 静态图中，核心是在什么层面获取算子记录在图中。 `jit.script`是在`python源码`角度分析function的源码，将`python code`转换为`图`（torchscript IR格式）。由于是直接从python源码转的，因此有许多python语法无法完备支持，存在转换失败的可能性； 而j`it.trace`则是在`c++层面`获取算子，在`算子调用`时记录成图（torchscript IR格式）。由于需要下降到c++，因此需要输入一遍数据真正“执行”一遍。而获取的图具有确定性，即如果图存在分支，则`只能被trace记录其中一个分支`的计算路线； `TorchDynamo`是将python源码转变成`二进制`后，通过分析二进制源码，获取分析的算子和图。最新的PyTorch1.14（2.0）中将其作为torch.compile()的主要路线。 `torch.fx`是在`python层面算子调用`时记录的算子，输出是fx IR格式的图。 `LazyTensor`则是在`c++算子调用`时记录成的图，该调用会截获正常算子的运行，在用户指定同步时再整体运行已积累的图。 __总结__：除了原生支持的动态图外，PyTorch在静态图方面做了大量尝试，由于Python语法的灵活性导致PyTorch在图模式方面设计困难，当前提出的诸多路线各有优劣，直到近期`PyTorch2.0`概念的提出才逐渐确定以`TorchDynamo`作为接下来静态图的主要路线。 本章仅简要介绍了各路线所处的层次，而动态图及静态图各路线的具体源码分析将安排在后续文章。