PyTorch源码解析(1)- 整体预览

    Hurray 3516次浏览 2条评论 5429字

#深度学习 #框架 #PyTorch #源码

> 第一篇文章节作为整体预览,一方面介绍代码的整体架构,一方面普及一些知识便于非ai开发人员也能快速了解相关背景。后续文章将针对不同的细节进行更深入的源码剖析。 > 本文已由公司账号公开发表在知乎:[https://zhuanlan.zhihu.com/p/598044604](https://zhuanlan.zhihu.com/p/598044604) [TOC] ## 1. 代码架构 ![](https://file.hurray0.com/uploads/menu/151/3782e05fcf9569e36bc8dbbbb6f00586.png) 本文截取的是`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包包含的内容: ![](https://file.hurray0.com/uploads/menu/151/f29ae7888b2b505e37804b70fbcc519e.jpeg) 对比下方`源码`的torch目录: ![](https://file.hurray0.com/uploads/menu/151/d0db5ef9fb86eb60b9aece9233988af5.jpeg) 可以发现torch的wheel包安装的内容基本上就是torch目录下的python内容。而c++内容(csrc目录)并没有被复制过来,而是以编译好的动态库文件(\_C.cpython-*.so)代替。 在torch目录下的`__init__.py`文件中,可以看到将\_C(动态库)导入的地方。 ![](https://file.hurray0.com/uploads/menu/151/914b4501f7df9d7dc4711dab406f684f.jpeg) ### 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++项目编译。 ![](https://file.hurray0.com/uploads/menu/151/38c808ac7e02ba609c0f5de9407b9753.jpeg) 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。 ![](https://file.hurray0.com/uploads/menu/151/fe8657e2ce3bf5a9bd93f1a43ce7f709.jpeg){:class="width_60"} 在\_C目录下\_\_init\_\_.pyi中,我们可以看到\_TensorBase的类。.pyi文件是存根文件(stub file),表示公共接口。其中有各种成员函数的定义,但是没有具体的实现。 在文件的注释中,我们可以看到其`具体实现`是在python_variable.cpp中。 ![](https://file.hurray0.com/uploads/menu/151/e773b76797517ac473627dac002fe28d.jpeg){:class="width_60"} 在python\_variable.cpp文件中,我们可以看到PyTorch实际是用了`pybind`,将c++和Python进行交互的。 ![](https://file.hurray0.com/uploads/menu/151/834b41b43bb70d84cd867041a1f517fc.jpeg){:class="width_60"} 我们这里以Tensor.dtype为例,它的实现在THPVariable\_properties中。 找到这个结构体的实现,可以看到dtype具体绑定了THPVariable\_dtype()这个函数。因此在Python中执行tensor.dtype时,实际运行的就是这个c++函数。 ![](https://file.hurray0.com/uploads/menu/151/7641994ea6f5cca5e829276c76275f65.jpeg){:class="width_60"} ![](https://file.hurray0.com/uploads/menu/151/dc56e3f22f7676a0f6e3b69c7db1c3ad.jpeg){: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()`来执行。 其整个过程其实是将计算过程构成了一张`计算图`,然后运行这个图的根节点。这样先构成图,再运行图的方式我们称为`静态图`或者`图模式`。 ![](https://file.hurray0.com/uploads/menu/151/5cd2288de851978fe7394714ef9642c6.jpeg){:class="width_60"} 而在PyTorch中,我们可以在计算的任意步骤`直接输出结果`。 原因是,PyTorch每一条语句是同步执行的,即每一条语句都是一个(或多个)`算子`,被`调用时实时执行`。这种实时执行算子的方式我们称为`动态图`或`算子模式`。 ![](https://file.hurray0.com/uploads/menu/151/7a0a39d5571c990cc59c988194fbf73a.jpeg){:class="width_60"} __优缺点__ `动态图`的优点显而易见,可以兼容Python式的编程风格,实时打印编程结果,用户友好性做到最佳; `静态图`则在性能方面有一定优势,即在整个图执行前,可以将整张图进行编译优化,通过融合等策略改变图结构,从而实现较好的性能。 PyTorch原生支持动态图,但是也在静态图方面做了诸多尝试。我们将在下面一节简单介绍这些图模式。 ### 2. PyTorch图原理概述 本节简要讲解一下PyTorch动态图以及几种`静态图`实现的`原理`,这里主要是简要总结,不涉及源码。 后续将有其他的文章更详细地分析PyTorch各路径的源码及其实现。 PyTorch除了有原生的动态图之外,当前也在静态图方面有诸多路线尝试:分别有`torchscript`(`jit.script`、`jit.trace`)、`TorchDynamo`、`torch.fx`、`LazyTensor`。 这里用一幅图总结了各路线所处的层次: ![](https://file.hurray0.com/uploads/menu/151/39d0ee608f9d1d64478d7ec5e46ae709.jpeg) 首先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`作为接下来静态图的主要路线。 本章仅简要介绍了各路线所处的层次,而动态图及静态图各路线的具体源码分析将安排在后续文章。

最后修改: