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

分享自己在学习 PyTorch 源码时阅读过的资料。本文重点关注阅读 PyTorch 源码的经验和 PyTorch 的代码结构。因为 PyTorch 不同版本的源码实现有所不同，所以笔者在整理资料时尽可能按版本号升序，版本号见标题前[]。最新版本的源码实现还请查看 PyTorch 仓库。更多内容请参考：

Ubuntu 22.04 LTS 源码编译安装 PyTorchpytorch/CONTRIBUTING.md 机翻PyTorch 源码学习PyTorch 源码学习：从 Tensor 到 Storage-CSDN博客

文章目录 [1.0.1] Oldpan：Pytorch源码编译简明指南核心文件夹third_partytools [1.7.1] 孙港丶丶：查看 Pytorch 源码[1.10.2] jhang的回答​：如何阅读pytorch框架的源码？[2.0.0]吾乃阿尔法：一点 PyTorch 源码阅读心得[2.1.1]kiddyjinjin：pytorch 源码解读进阶版 - 总述[unknown] 小飞的回答​：如何阅读pytorch框架的源码？[unknown] 小飞​：PyTorch源码学习系列 - 1.初识我理解的PyTorch架构PyTorch目录结构 [unknown] clay001​：pytorch源码阅读待更新……

[1.0.1] Oldpan：Pytorch源码编译简明指南

具体内容见原文

以下是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。

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

[1.7.1] 孙港丶丶：查看 Pytorch 源码

具体内容见原文

该图来自：PyTorch internals : ezyang’s blog

“torch/” 中的代码文件一般是pytorch python类的接口信息，其内容可以直接在编辑器中通过查看定义得到，但其只包括了接口的参数信息和注释，注释与官方文档中收到的内容相同；“torch/csrc/” 中含有python和c++衔接的代码，很多为编译时生成；“aten/src/ATen/” 中包括了torch中许多操作的C++或C实现，是我们查看pytorch许多函数实现需要关注的区域；“c10/” 中为torch最基本特性实现的部分，如张量的定义、数据类型的定义等。

因此，为查看torch函数的具体实现，需要查看"aten/src/ATen/" 部分的内核代码。

其中"native/“为C++风格的实现，”/TH"等主要为C风格的实现，Edwards 说开发人员在努力将"/TH"等中的操作迁移到"native/"中，至少在我使用的pytorch1.7.1中，"native/"覆盖了大部分的张量操作，如卷积、激活和损失等等。

如何查看"aten/src/ATen/native"中的函数操作？

首先可以在"aten/src/ATen/native/native_functions.yaml"这个注册文件中查看函数的框架，可以确定这个函数是否在"native/"中实现，这里的框架主要是用于代码生成。直接搜索大法，搜索文件夹中所有文件的内容，找到你想要的函数。\ [1.10.2] jhang的回答​：如何阅读pytorch框架的源码？

具体内容见原文

要想阅读 pytorch 框架的源码，最起码得先了解框架的调用栈，以 pytorch v1.10.2的 torch.nn.conv2d 为例，介绍一下python api 接口到底层 c++ 接口的调用依赖关系：

Python API:

torch.nn.Conv2d: ---> class Conv2d, (torch/nn/modules/conv.py) ---> F.conv2d, (torch/nn/functional.py) ---> def conv2d(input: Tensor, ...) -> Tensor: ..., (torch/_C/_VariableFunctions.pyi)

C++ API:

---> at::Tensor conv2d(...), (aten/src/ATen/native/Convolution.cpp: 579) ---> at::convolution(...), (aten/src/ATen/native/Convolution.cpp: 586) ---> at::Tensor convolution(...), (aten/src/ATen/native/Convolution.cpp：743) ---> at::_convolution(...), (aten/src/ATen/native/Convolution.cpp：754) ---> at::Tensor _convolution(...), (aten/src/ATen/native/Convolution.cpp：769) ---> at::cudnn_convolution(...), (aten/src/ATen/native/Convolution.cpp：860) ---> Tensor cudnn_convolution(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：265) ---> cudnn_convolution_forward(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：273) ---> Tensor cudnn_convolution_forward(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：221) ---> raw_cudnn_convolution_forward_out(...), (aten/src/ATen/native/cudnn/ConvShared.cpp：258) ---> void raw_cudnn_convolution_forward_out(...), (aten/src/ATen/native/cudnn/Conv_v7.cpp：669) ---> split_batch_dim_to_32bit_out(...,raw_cudnn_convolution_forward_out_32bit), (aten/src/ATen/native/cudnn/Conv_v7.cpp：673) ---> void raw_cudnn_convolution_forward_out_32bit(...), (aten/src/ATen/native/cudnn/Conv_v7.cpp：625) ---> cudnnConvolutionForward(...), (aten/src/ATen/native/cudnn/Conv_v7.cpp：658)

python 端接口比较容易查找，对于 c++ api 接口的查找给与一些建议：

常用 op c++ 接口都位于aten/src/ATen/native目录，而且该目录下有个native_functions.yaml, 记录了c++接口的调用依赖关系kernel 名搜索：通常我们也会通过 nvprof 查看调用 kernel 名逆向查找 c++调用接口关键字搜索：以conv2d为例, 我们也会去搜索 Tensor conv2d 关键字来查找，方便将 python 端的最后一个接口与 c++ 端的第一个接口对应起来；先验知识搜索，还需要拥有一些cudnn api的先验知识，比如我知道conv2d一定会调用 cudnnConvolutionForward 等接口, 然后搜该接口所在位置，然后逆向搜索 c++ api 接口即可c++ 接口的搜索一般只需要找到离 cuda kernel 或 cudnn/cublas API 最近的位置即可(所以方法 2，4是我们最常用的手段)，从 python 接口到 c++ 接口的中间依赖关系没必要深究，因为其中做了很多封装，经常依赖 vscode 跳转时会断了依赖关系；推荐使用vscode开发工具做好配置，方便c++/python 接口的跳转，有助于理解调用过程 [2.0.0]吾乃阿尔法：一点 PyTorch 源码阅读心得

具体内容见原文

PyTorch 的代码对初学者来说并不好读，原因在于：

有很多 C++ 代码是在编译过程中生成的；PyTorch 前端充分利用了 Python 的动态特性，比如 Python 函数可能经过 N 个装饰器修饰；混合 Python/C++ 调用，比如从 Python 调 C++ 调 Python 再调 C++ 在 PyTorch 中是比较常见的事情；

以下是我读 PyTorch 源代码的一些经验，仅供参考：

明确目标: PyTorch代码很多，建议每次只读一个专题或者从一个特定的问题出发，比如 PyTorch AMP 是怎么实现的；

把握全局: 一上来直接读代码会有很多障碍，很多术语不明所以，先通读这个专题下官方的教程、文档，以及一些写的好的第三方博客，确保自己对这个专题下的内容有初步的认知。以下是一些初步了解 PyTorch 特定专题内容的比较好的资源:

PyTorch 教程PyTorch 文档PyTorch 博客PyTorch 案例PyTorch 论坛PyTorch Youtube 频道

Debug Build: 一定要 build debug 版的 PyTorch，并保留在编译过程中生成的源代码，否则很难找到 PyTorch 函数调用栈中一些函数的来源；

此处省略细节，具体内容见原文

静态读代码: 有了完整的 PyTorch 源代码之后就可以开始读了，网上有很多 VSCode 教程，设置好 VSCode 的 Python 和 C++ 插件，方便在函数之间跳转，可以解决一大部分的函数跳转；

动态读代码: 静态读代码的问题是常常搞不清函数的执行流程，此时在运行过程中动态读执行过的代码就很有帮助，善用 gdb 和 pdb 可以有效辅助读源代码。

此处省略细节，具体内容见原文

充分利用源代码中的日志、debug 选项、测试用例: 很多 PyTorch 模块都包含了丰富的日志和 debug 开关，这些日志和用于调试的消息可以帮助我们理解 PyTorch 的执行流程。除此之外，PyTorch 中包含了大量的测试用例，如果单纯看源代码无法理解程序的逻辑，看看其对应的测试用例可以帮助我们理解程序在做什么。

此处省略细节，具体内容见原文

及时求助: 如果经过上面的流程还无法了解某些代码的逻辑，要及时向社区求助，避免浪费过多时间。

学什么: 明确源代码中哪些东西值得我们学习和借鉴，读源代码时要特别注意这些方面，比如:

特定模块/功能的实现原理;用到的算法;一些 coding 技巧;

知行合一: You can’t understand it until you change it. 读源代码不是最终目的，充分利用从代码中获取的认知才是。有效的输出可以加深我们对代码的理解，一些可以参考的输出方式:

写一篇源码剖析的博客;简化自己对源代码的认识，分享给其他人;修改源代码，改进或添加一些功能，给 PyTorch 提交 PR;亲手实现相同功能，写精简版的代码复现核心逻辑;

每一个读源代码的人都是不甘平凡的人，祝大家在这个“痛并快乐着”的过程中成长得更快、更多。

[2.1.1]kiddyjinjin：pytorch 源码解读进阶版 - 总述

具体内容见原文

pytorch官方给出的参考代码结构：pytorch/pytorch

源码解读的代码版本基于v2.1.1 版本：GitHub - pytorch/pytorch at v2.1.1

其中比较核心的4个目录简单介绍如下：

torch：python 代码部分， the frontend, 我们在代码中引入并使用的Python 模块(modules)，基本都在这里torch/csrc: python bindings, 这部分C++代码实现了所谓的PyTorch前端(the frontend of PyTorch)。具体来说，这一部分主要桥接了Python逻辑的C++的实现，和一些PyTorch中非常重要的部分，比如自动微分引擎(autograd engine)和JIT编译器(JIT compiler)。aten/src/ATen: 是“A Tensor Library”的缩写，是一个C++库实现了Tensor的各种operations。ATen 内对operators的实现分成两类，一种是现代的C++实现版本（native），另一种是老旧的C实现版本（legacy）。c10: 是一个来自于Caffe2 和 ATen的双关语(Caffe 10)，其中包含了PyTorch的核心抽象，比如 Tensor、Device、Allocator、Storage 等数据结构的实际实现部分。

进一步对c10/core 的代码结构进一步介绍如下：

Allocator.h/cpp: 提供memory管理的抽象类定义和实现。CPUAllocator.h/cpp: set/get CPU Allocator, 指定DefaultCPUAllocator。Device.h/cpp: 提供底层硬件的抽象类定义和实现。 struct C10_API Device final 表示张量所在的计算设备, 设备由类型唯一标识（cpu or cuda gpu， cuda 的话有index） DeviceGuard.h : RAII guardDeviceType.h 定义了21 种 DeviceType，包括我们常见的CPU，CUDA，XLA，HIP 等，应该会是逐渐增加的状态。DispatchKey.h 定义了40+种 dispatcherkey set，是个 uint8_t，越高位优先级越高，用来给op路由，当调用一个op的时候，pytorch中的dispatch 机制会结合 DispatchKey 以及当时的场景（在什么硬件上？是推理还是训练？）分配到特定的执行算子上，这个机制保证了pytorch可以非常灵活的在不同的硬件和场景下切换。后面会详细讲讲，比较有意思。DispatchKeySet.hTensorOptions.h/cpp: 提供创建张量时的选项（如设备类型、数据类型等）的定义和相关操作。TensorImpl.h/cpp: 提供张量的具体实现。 其他：Backend.h （老旧版本的 Device.h/cpp， layout 信息），提供多种新旧互相转换接口 [unknown] 小飞的回答​：如何阅读pytorch框架的源码？

具体内容见原文

分享下我个人学习经验

首先自己一定要编译源码，学会用GDB去调试代码。任何源码学习的第一步永远是将源码编译成功并运行起来。其次切记不要去看 setup.py 文件，不要去看 torch/__init__.py 文件，不要去看C++编译配置文件。如果你没有丰富的相关工程开发经验，从你打开这些文件起，恭喜你，你已经被成功劝退了。初学者就应该老老实实先学习项目的设计思想和架构，忽略工程配置。了解下Python语言如何与C/C++语言绑定。PyTorch底层核心代码都是用C/C++编写，如果你仅仅想学习的是Python部分的源码，那本回答就到此结束了，根据你python里的模块用到哪学到哪就行了。如果你想深入学习C/C++源码，这是一道必须迈过的坎。不理解的话你对PyTorch的认识永远是Python和C/C++两个孤零零的世界，无法联系起来。先学习Tensor。Tensor是贯穿整个PyTorch最基本的数据结构，无论何时你都会遇到它。了解算子。通过最简单的Tensor.Add 方法学习下Tensor函数是如何设计的。PyTorch里操作Tensor的函数被称之为算子，算子包含前向和后向计算。这里不要求完全弄懂算子，只需要了解相关流程就行。学习Autograd。先了解计算图的概念，学习下PyTorch是如何创建和维护计算图。如何反向传播自动求解梯度更新权重。

学到这里你基本上对PyTorch C/C++底层架构有所了解，下面就是根据自己喜爱的方向去学习。

如果你想研究深度学习模型，去看Python torch.nn包里面的python源码就行。如果你想研究算子实现，还记得我们前面学的Add方法了，找到你想了解的算子去深入，如果不关心速度，看下CPU代码就行，如果想学习如何优化就去了解cuda编程。如果想学习分布式，就去研究下分布式相关的code。 其实从上面的流程看，只要你前5步都完成了，你自然而然就知道该如何去学习PyTorch框架源码。现在PyTorch源码很大，想在有限时间内一下子全看完不现实，必然要结合自己的实际情况有所侧重。 [unknown] 小飞​：PyTorch源码学习系列 - 1.初识

具体内容见原文

PyTorch本身是一个Python扩展包，按照官方说法它主要具有以下两种特色：

支持GPU加速的张量（Tensor）计算在一个类似磁带（前向和反向）的梯度自动计算（Autograd）系统上搭建深度神经网络

Tensor其实本质上就是一个多维数组。在数学上单个数据我们称之为标量，一维数据我们称之为向量，二维数据我们称之为矩阵。

如果PyTorch仅是支持GPU加速的Tensor计算框架，那它也就是NumPy的替代品而已。其最核心的功能就是Autograd系统，目前深度学习就是基于梯度反向传播理论来达到网络的自我训练。PyTorch的Autograd系统是创建了一个动态的计算图，用户只需要关注前向计算网络搭建，PyTorch会自动根据动态计算图去反向计算梯度并更新网络权重。

在设计之初PyTorch并没打算仅成为一个绑定C++框架的Python包，它紧密地将C++框架集成到python中。你可以在使用PyTorch的同时也结合使用NumPy/SciPy/scikit-learn这些优秀的科学计算包，同时你也可以用Python编写你自己的神经网络层（PyTorch神经网络相关的代码基本上都是Python编写）。

正如前面所说，PyTorch的计算图是动态的，当你写下一行Python代码的时候你就可以直接运行得到结果。对用户来说，PyTorch的世界就是一个简单的单线程同步非阻塞世界，你可以得到实时反馈，这对于刚开始接触深度学习的新手来说是非常贴心的功能，可以帮忙新手快速入门。

我理解的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中很多相似源码都是脚本通过模板自动生成的，这个文件夹下面就放着自动生成代码的脚本 [unknown] clay001​：pytorch源码阅读

具体内容见原文

pytorch代码主要由C10（Caffe Tensor Library，最基础的Tenor库代码），ATen（A Tensor library for C++11，基于C10），torch三大部分组成。

torch.nn中包含各种神经网络层，激活函数，损失函数的类，使用时需要先创建对象。

torch.nn.functional中的接口可以直接调用函数而不用创建对象。

算子配置文件：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的形式存在。

待更新……