当前位置： 首页 > news >正文

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

news 2025/2/22 2:09:57

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

  • Ubuntu 22.04 LTS 源码编译安装 PyTorch
  • pytorch/CONTRIBUTING.md 机翻
  • 深度学习框架与静态/动态计算图【笔记】
  • PyTorch 源码学习
  • PyTorch 源码学习：从 Tensor 到 Storage-CSDN博客

文章目录

  • [1.0.1] Oldpan：Pytorch源码编译简明指南
    • 核心文件夹
    • third_party
    • tools
  • [1.4.0] 技术瘾君子1573：深度学习Pytorch框架扩展知识详解
  • [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源码包展开的目录结构(只展示了主要的一些文件夹)，其中主要的源码都在以下展示的文件夹中：

《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.4.0] 技术瘾君子1573：深度学习Pytorch框架扩展知识详解

具体内容见原文

├── android: 用于构建android应用程序
├── aten: A Tensor Library，pytorch的C++ tensor 库,实现了tensor的运算 (不支持自动梯度)，大部分op都可以在此找到，Aten的op实现分为两种;1. native op：op的现代C++实现 2. legacy op(TH, THC, THCUNN):原torch遗留下来的op的传统C语言实现,不要在里面添加新的op,这部分正在慢慢转移到native中
│  └── conda:包含了conda构建使用的脚本以及配置文件
│  └── src：op的实现
│     └──ATen：op的现代C++实现
│       └──core：Aten的核心功能，目前正在迁移到c10文件夹中
│       └──native：op的现代实现，在这里添加你想要添加的新的op
│         └──cpu：不是op的实际cpu实现，但是包含了特殊指令集的实现，比如AVX
│         └──cuda：op的cuda实现
│         └──sparse：COO的稀疏tensor的op的cpu和cuda的实现
│         └──mkl mkldnn miopen cudnn：绑定到具体后端的op的实现
│     └──TH：TorcH,CPU tensor的实现。TH,THC,THCUNN是原torch遗留下来的op的C语言实现
│       └──generic：TH,THC,THCUNN三者的generic中包含了操作的实际实现，这些文件会被编译N次
│       └──…
│     └──THC：GPU tensor的实现
│       └──generic
│       └──…
│     └──THCUNN：GPU版底层神经网络实现
│       └──generic
│     └──THH：THC在hip上的实现
│     └──THHUNN：THCUNN在hip上的实现
│  └── tools:包含了一些测试脚本
├── benchmarks:包含了一些pytorch特性的benchmark测试
├── binaries
├── c10:核心Tensor库，从ATen/core迁移过来，只包含最基本的功能，可以运行在服务端和移动端，适用于对二进制文件大小有要求的环境中
│  └── benchmark
│  └── core：CPU的核心库
│  └── cuda：cuda的核心库
│  └── hip：HIP的构建文件,该文件夹不包含实际的文件，这些文件通过tools/amd_build中的HIPIFY脚本拷贝过来
│  └── macros:包含了一些自定义的宏
│  └── test：测试程序
│  └── util：通用的工具库
├── caffe2: caffe2
│  └── core：caffe2的核心文件，包括tensor,workspace,blobs等
│  └── operators：caffe2的op
│  └── python：caffe2的python绑定
│  └── …
├── cmake:用于构建pytorch的cmake文件
├── docker:用于在docker中使用pytorch
├── docs:pytorch相关的说明文档
├── ios:pytroch的ios模块，用于构建ios应用程序
├── modules
├── scripts:包含了一些有用的工具,比如build_android.sh用于创建android上的pytorch库
├── submodules
├── test:pytorch的python前端的python单元测试
│  └── test_torch.py:pytorch功能的基本测试
│  └── test_autograd.py:非神经网络自动微分的测试
│  └── test_nn.py :神经网络op以及自动微分的测试
│  └── test_jit.py: JIT 编译器和TorchScript的测试
│  └──…
│  └──cpp:pytorch的C++前端的C++单元测试
│  └──onnx:onnx输出功能的测试，同时使用pytorch和caffe2
├── third_party:pytorch依赖的第三方库
├── tools:pytorch库的代码生成脚本
│  └── autograd: 生成自微分相关函数的工具
│  └──…
│  └──shared
├── torch:实际的pytorch库。包含了我们平时import的python模块，所有不在csrc文件夹中的是一个python模块，符合pytorch的python前端模块结构
│  └── csrc: pytorch库的C++实现，包含1. C++前端的实现 2.绑定代码(python bindings)。csrc中每个独立文件与python模块对应。python绑定文件的规范列表请查看setup.py,通常他们以python_为前缀
│     └──api:pytorch的C++前端接口
│     └──autograd：自动微分反向模式的实现
│     └──cuda
│     └──distributed：pytorch的分布式训练
│     └──generic
│     └──jit: TorchScript JIT的编译器和前端
│     └──…
│     └──utils
│     └──Module.cpp:构建Python扩展模块
│  └── nn:torch.nn模块
│  └── optim:torch.optim模块
│  └── distributed:torch. distributed模块

总体结构上，C/C++ 后端用aten库作为上层封装再通过torch/csrc中的部分胶水代码接入 Python (部分为 pybind11)。

实际上，Pytorch 中重点为下面 4 个目录：

  • torch/
  • torch/csrc/
  • aten/src/ATen/
  • c10/

[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"中的函数操作？

  1. 首先可以在"aten/src/ATen/native/native_functions.yaml"这个注册文件中查看函数的框架，可以确定这个函数是否在"native/"中实现，这里的框架主要是用于代码生成。
  2. 直接搜索大法，搜索文件夹中所有文件的内容，找到你想要的函数。\

[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 接口的查找给与一些建议：

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

[2.0.0]吾乃阿尔法：一点 PyTorch 源码阅读心得

具体内容见原文

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  9. 知行合一: 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 guard
  • DeviceType.h 定义了21 种 DeviceType，包括我们常见的CPU，CUDA，XLA，HIP 等，应该会是逐渐增加的状态。
  • DispatchKey.h 定义了40+种 dispatcherkey set，是个 uint8_t，越高位优先级越高，用来给op路由，当调用一个op的时候，pytorch中的dispatch 机制会结合 DispatchKey 以及当时的场景（在什么硬件上？是推理还是训练？）分配到特定的执行算子上，这个机制保证了pytorch可以非常灵活的在不同的硬件和场景下切换。后面会详细讲讲，比较有意思。
  • DispatchKeySet.h
  • TensorOptions.h/cpp: 提供创建张量时的选项（如设备类型、数据类型等）的定义和相关操作。
  • TensorImpl.h/cpp: 提供张量的具体实现。
    其他：
  • Backend.h （老旧版本的 Device.h/cpp， layout 信息），提供多种新旧互相转换接口

[unknown] 小飞的回答​：如何阅读pytorch框架的源码？

具体内容见原文

分享下我个人学习经验

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

待更新……


http://www.mrgr.cn/news/91687.html

相关文章：

  • 嵌入式音视频开发（三）直播协议及编码器
  • 【Java】泛型与集合篇 —— Set 接口
  • 前端常见面试题-2025
  • C语言——时基
  • Linux-----进程（多任务）
  • C#发送邮件
  • 基于正则化密集连接金字塔网络的显著实例分割
  • mysql总结
  • Day6 25/2/19 WED
  • Windows 启动 SSH 服务报错 1067
  • Compose 常用UI组件
  • PVE使用一个物理网卡采用VLAN为管理IP和VM分配网络的问题
  • springboot-ffmpeg-m3u8-convertor nplayer视频播放弹幕 artplayer视频弹幕
  • 【SQL】多表查询案例
  • OpenResty
  • [数据结构]顺序表详解
  • 驱动开发、移植
  • 动态订阅kafka mq实现(消费者组动态上下线)
  • 【3.2JavaScript】JavaScript语法基础
  • git-提交时间和作者时间的区别