PyTorch 源码学习：阅读经验 & 代码结构

PyTorch 源码学习：阅读经验 & 代码结构

引用

CSDN

1.

https://blog.csdn.net/weixin_43254181/article/details/145550497

本文是一篇关于PyTorch源码学习的详细指南，内容非常丰富，涵盖了多个版本的PyTorch源码分析和阅读经验分享。文章从多个角度介绍了PyTorch的代码结构、核心组件、以及如何有效地阅读和理解源码。

PyTorch 源码学习：阅读经验 & 代码结构

一些资料和文档

入门书：

• 深度学习入门 (豆瓣)
• 深度学习入门2 (豆瓣)

官方资源：

• 官方网站
• 源码仓库
• 中文文档
• 开发者文档
• 贡献者文档

Brief PyTorch Architecture

PyTorch框架由使用C++、Python和CUDA编程的张量库组成。这些库存储并操作张量，张量是多维的矩形数字数组。应用程序必须导入torch库才能使用PyTorch的功能。图1展示了PyTorch架构的简化视图。

核心文件夹

核心文件夹主要是c10、aten、torch、caffe2。

为什么将c10放到最前面呢？
因为官方已经表明c10目录是最重要的源代码文件夹，也就是几乎所有的源代码都与这里的代码有关系，比如我们的类型定义，Pytorch最重要的Tensor的内存分配方式等等，官方也说到了，之后会慢慢将Aten中的代码移至这个文件夹，也就是说这个文件夹将包含Pytorch中最核心的代码。

而Aten文件夹则包含了一些实现了Tensor的底层(和c10类似)，也包括了很多的层前向代码和后向实现的代码(例如卷积层的前向和后向操作代码)，包括CPU和GPU端，总之都是C++的核心操作代码。

torch文件夹也同样重要，其中主要包含了一些稍微高层些的操作函数，例如torch.ones等，有C++和Python端，也包括了Python核心代码和包装代码，如果我们使用python版Pytorch的话，与这些代码接触就比较密切了。

而Caffe2则不用多说，caffe2则主要针对移动端设计了很多优化后的运算代码，模型融合、模型量化等等的代码，其后端有QNNPACK等一些针对移动端的底层运算库(有开发人员说GLOW也在caffe2后端考虑之内)。

third_party

Pytorch毕竟是大型的深度学习库，所以需要的依赖库也是有很多的，其中有很多我们耳熟能详的数值计算库(eigen、gemmlowp)、模型转换库(onnx、onnx-tensorrt)、并行训练库(gloo、nccl)、自家的底层端实现库(QNNPACK)以及绑定python端的pybind11等一系列所依赖的库。

当然还有很多库这里就不一一介绍了，总之，我们在编译的时候，Pytorch的编译代码会根据我们的设置在编译的时候，自动判断当前系统中是否存在需要的第三方库。如果不存在则使用这里的第三方库(直接编译并使用第三方库的diamante)，这也是为什么我们需要执行

git submodule update --init --recursive

来下载所依赖第三库源码的原因。

tools

tools这个文件夹中的内容到底是做什么的，简单看一下官方的介绍：

This folder contains a number of scripts which are used as part of the PyTorch build process. This directory also doubles as a Python module hierarchy.

这个文件夹包含许多用作PyTorch构建过程一部分的脚本。该目录还兼作Python模块层次结构。

其中包含了一些脚本生成代码工具(利用python)、用于编译一些组件的脚本和代码，还有一些开发人员需要的工具、以及AMD显卡帮助编译代码和一些特殊情况需要使用的工具等。在我们编译Pytorch源码的过程中会使用到这个文件夹中的代码。

有一点需要说明，那就是Pytorch利用了很多的代码生成，例如操作层函数的头文件NativeFunction.h等，所以tools中的代码生成脚本还是比较重要的。

提一个可能会使用到的脚本build_pytorch_libs.sh，这个脚本是用来编译libtorch库的，libtorch就是不需要python包装的使用C++的Pytorch库，方便于部署阶段使用。

关于tools中的文件就不具体介绍了，大家可以看一下其中的readme。

其他的文件夹就不多说了，相对上面的来说并不是很重要。

我理解的PyTorch架构

根据自己的理解简单画了下PyTorch的架构图，隐藏了部分细节。本系列也将主要围绕着这张架构图去学习PyTorch的具体实现。

一共将PyTorch分成了四层，分别是

• 应用层（Python）。这应该是大家最熟悉的层，主要涉及到张量，Autograd以及神经网络。该层所有的源码都是由Python编写，这也符合前面所说的PyTorch设计思想-——将C++框架集成到Python里
• 实现接口层（C++）。该层的主要功能我认为有两个：
• Python 扩展。通过Python提供的C API将Python应用层与C++实现层绑定起来，使用户在享受Python语言提供的便捷优势时也可以同时享受到C++语言提供的性能优势
• Autograd系统实现。 PyTorch并没有在实现层中实现Autograd系统。在此层中PyTorch定义了动态有向图的基本组件Node和Edge，以及在此基础上封装了Function类和Engine类来实现Autograd
• 实现层（C++）。该层是PyTorch的核心层，定义了PyTorch运行过程中的核心库，包括Tensor的具体实现，算子实现（前向与后向运算）以及动态调度系统（Tensor的布局，硬件设备，数据类型）。Storage类主要是针对不同硬件数据存储的一种抽象。
• 硬件接口层。该层主要是硬件厂商基于自家硬件推出的运算接口。

PyTorch目录结构

PyTorch的源码托管于GitHub平台，其目前的代码量已经非常巨大。新手第一次接触的时候往往会因此被劝退，但其实里面很多文件和我们学习PyTorch源码并没有太多的直接关系，所以我们第一步就是要理清目录结构，专注于我们需要学习的内容。

• torch：我们“import torch”后最熟悉的PyTorch库。所有非csrc文件夹下的内容都是标准的Python模块，对应我们架构图中的应用层
• csrc：该目录下都是C++源码，Python绑定C++的相关code都在这个目录里面，同时也包含了对PyTorch核心库的一些封装，对应我们架构图中的实现接口层
• csrc/autograd：梯度自动计算系统的C++实现
• autograd：梯度自动计算系统的Python前端源码，包含torch中支持的所有自动求导算子
• nn：建立在autograd系统上的神经网络库，包含了深度学习中常用的一些基础神经网络层。
• optim：机器学习中用到的优化算法库
• aten：“a tensor library”的缩写，对应我们结构图中的实现层。从名字上也能知道，这个库设计之初主要是为Tensor服务。因为在实现接口层下面，所以这里的Tensor并不支持autograd
• src/Aten/core：aten的核心基础库。目前这个库里面的代码正在逐渐地迁移到c10目录下面
• src/Aten/native：PyTorch的算子库，这个目录下面的算子都是CPU的算子。对于一些专门CPU指令优化的算子会在子目录里面
• src/Aten/native/cuda：cuda算子实现
• c10：“caffe2 aten”的缩写，PyTorch的核心库，支持服务端和移动端。
• tools：PyTorch中很多相似源码都是脚本通过模板自动生成的，这个文件夹下面就放着自动生成代码的脚本

如何阅读PyTorch源码

阅读方法

面对这么多的代码和文件，一下子肯定不知所措，尤其是阅读新模块的时候，我首先会尝试找到该模块的说明，通过README.md或前人的博客或API文档了解下该模块大概功能和结构，然后整体（粗略）浏览一遍该模块的代码，对每个文件里的代码是做什么的有个大致概念，最后再根据自己的理解选择性地进行精读。

特定功能分析

1. GPU内存管理：

• TensorFlow BFC算法的分析
• CUDACachingAllocator的深入分析
• expandable_segments的初步探索

2. 自动微分（Autograd）：

• Dispatch机制的详细解析
• Operators的实现细节

阅读建议

1. 明确目标：每次只读一个专题或者从一个特定的问题出发，比如PyTorch AMP的实现。
2. 把握全局：先通读官方教程、文档以及第三方博客，确保对专题内容有初步认知。
3. Debug Build：一定要build debug版的PyTorch，并保留编译过程中生成的源代码。
4. 静态读代码：设置好VSCode的Python和C++插件，方便函数间跳转。
5. 动态读代码：善用gdb和pdb辅助读源代码。
6. 充分利用源代码中的日志、debug选项、测试用例：很多PyTorch模块都包含了丰富的日志和debug开关，这些可以帮助理解执行流程。
7. 及时求助：如果无法理解某些代码逻辑，要及时向社区求助。
8. 知行合一：读源代码不是最终目的，要充分利用从代码中获取的认知，比如写一篇源码剖析的博客、简化认识分享给他人、修改源代码提交PR等。

算子配置与实现

算子配置文件

pytorch/aten/src/Aten/native/native_functions.yaml中有关于各个算子的说明，同级目录下有这些算子的实现。每个算子有一些字段：func，variants，dispatch。

• func字段：表示算子的名称和输入输出类型
• variants字段：表示高级方法的使用方式，
• function表示可以用torch.xxx()的方式调用，
• method表示在tensor a上，用a.xxx()的方式调用
• dispatch字段：分发的设备类型，CPU，CUDA等等

算子通常成对出现（正向和反向）。python_module: nn表示将该方法自动生成到torch.nn.functional模块中，可以通过torch.nn.functional.leaky_relu的方式来调用。

反向算子配置文件

在tools/autograd/derivatives.yaml中添加算子和反向算子的对应关系。

算子实现文件

• aten/src/Aten/native/目录下的h和cpp文件。有些会按照功能实现在一起，比如Activation.cpp中包含了多个算子。在这些cpp实现中，用到了封装后的函数，会再往里调一层。
• aten/src/ATen/native/cpu/目录下，通常以xxx_kernel.cpp的形式存在。

总结

PyTorch的源码虽然庞大，但通过本文的结构化分析和阅读建议，相信读者能够更有效地理解和掌握其核心原理。无论是对深度学习框架感兴趣的开发者，还是希望深入理解PyTorch实现的AI工程师，本文都提供了一个全面且实用的指南。