简介这份PyTorch深度学习源码包面向初学者与中级开发者集中整理了深度学习实践中的常用代码与示例覆盖数据预处理、模型构建、训练与评估等环节有助于解决入门时源码分散、难以跑通的问题。压缩包共含140个文件大小29.33MB其中以45个ipynb交互式笔记、37个py脚本为主体另有14张png图片、9个Markdown说明文档以及sample样例、配置文件等目录结构清晰便于按需检索。目前已有1943人学习下载受到同类学习者验证。源码大多可直接运行若遇环境不兼容可依报错信息上网排查借助notebook分步演示、py脚本模块化实现与配套图表读者能快速复现经典模型、理解PyTorch编程范式可作为课程设计、项目实战或论文复现的起步参考。1. 深度学习PyTorch源码一份值得从第一行读到底的框架代码如果只挑一个深度学习框架去读源码PyTorch 是性价比最高的选择至少我拆完之后是这么认为的。你在训练脚本里写下的每一行 tensor 运算都不只是 Python 层面的一次方法调用它会越过 Python 和 C 的边界落到 ATen 张量库和 autograd 引擎里执行最终才变成你在控制台看到的那个 loss。这份源码解析把整个走读路径拆成可执行的动作先看目录结构建立地图然后本地构建一个能单步调试的版本接着精读 autograd 和 nn.Module 的关键实现最后演示怎么注册一个自定义算子并验证梯度。适合已经被 bug 折磨过、想知道 loss.backward() 里那几毫秒到底在发生什么的调参工程师也适合需要自研算子、做模型导出或训练框架二次开发的工程师。读完你至少能回答一个问题当你写出 loss.backward() 时源码里是哪一行代码在真正干活。2. 架构设计源码目录地图与一次张量调用的完整落点读 PyTorch 源码最怕的不是代码量大而是只盯着 torch/ 目录下那几个 Python 文件当全部结果发现真正干活的地方在 C 边。这一章先铺地图再从一次最普通的张量加法出发把一层层调用关系拉直。看完之后你就知道为什么 Python 端那么“薄”C 端又在哪里拦下了绝大部分逻辑。2.1 从目录结构看分层Python 薄壳与 C 引擎拿到源码仓库后第一件事是展开根目录把它的结构分层记在脑子里。PyTorch 仓库不是一个大平层它至少可以切出三层对外暴露的 Python 包、pybind11 生成的绑定层、以及纯 C 实现的核心库。pytorch/ ├── torch/ │ ├── _C/ # pybind11 绑定层编译后生成 _C.so / _C.pyd │ ├── _tensor.py # Tensor 的 Python 方法入口 │ ├── _torch_docs.py # 独立维护的 docstring │ ├── nn/ # Module、Parameter、functional、utils │ ├── optim/ # SGD、Adam、AdamW、lr_scheduler │ ├── autograd/ # backward、Function、grad_mode 的 Python 侧 │ ├── utils/ # data、checkpoint、cpp_extension 等 │ └── __init__.py # 组装入口 ├── torch/csrc/ │ ├── aten/ # ATen 张量库的核心 │ ├── autograd/ # 自动求导引擎含 generated/ │ └── api/ # 对外的 C API 封装 ├── third_party/ # 第三方依赖如 pybind11、gloo └── test/ # 测试与回归样例先讲清楚几个容易误解的目录。torch/ 下的 Python 包是绝大多数用户直接接触的部分但它只是一个转发层比如 torch/_tensor.py 里的 Tensor.add 方法只做参数组装和协议分发真正执行加法的 CPU kernel 在 torch/csrc/aten/src/ATen/native/cpu/ 下面。torch/csrc 是 pybind11 绑定和自动求导引擎的所在地其中 _C 模块编译成扩展文件后会在 torch/init.py 里被 import整个 torch 命名空间里的函数如果找不到 Python 实现最终都会顺着 _C 走到这里。还有一个容易被忽略的文件是 native_functions.yaml它位于 aten/src/ATen/native/ 下是所有算子的 schema 声明文件代码生成器会读它生成绑定、分发和导数接口后面改自定义算子的时候你会回头来找它。我一般会按“入口文件 → schema 声明 → C kernel”的顺序去读任意一个算子而不是从 YAML 往前读因为从用户侧的代码往回追更容易建立直觉。这一章后面的链路演示也沿用了这个顺序。2.2 一次 a b 到底走了几层Python 方法到 CUDA Kernel用一个最普通的张量加法来把链路走通。先写一个最小脚本import torch a torch.randn(4, 4) b torch.randn(4, 4) c a b # 在这一行打断点 print(c.grad_fn) # None因为没有 requires_grad d (a b).sum() # 再看一次加法如果对 Python 层做断点调试比如在 IDE 或 pdb 里逐行进你会看到 a b 被解析到 torch/_tensor.py 里的 Tensor.add它会先做一次类型检查然后调用 _C._VariableFunctions.add 这个 pybind11 绑定函数。到这里 Python 层就结束了控制权从 Python 解释器切换到 C 扩展。进入 C 之后add 会先查找 Dispatcher 里的分发表Dispatcher 按 device 类型和 layout 去匹配 CPU 或者 CUDA 的实现CPU 上最终落在 aten/src/ATen/native/cpu/BinaryOpsKernel.cpp 里那个 kernel 函数CUDA 上则落到对应的 .cu 文件。匹配过程里的 key 包括 device、dtype、layouts 等PyTorch 里管这套机制叫 dispatch 体系。它保证了同样的 add API 在 CPU、GPU、MPS 不同后端上都能被调度到正确的实现而用户侧完全不需要感知。这一层在 Python 侧留了两个协议可以拦截torch_function和torch_dispatch。前者拦截的是 Python 层操作适合做 Tensor 子类后者拦截的是已经进入 C 分发层之前的操作适合做各种追踪和自定义 ExecutionMode。如果你要二次开发大多数“改算子行为”的入口都在这两个协议附近而不是直接去改 C 源码。这个判断用在很多项目里都适用先找协议再找 kernel。2.3 动态图在源码里的落点AutogradMeta 与 grad_fnPyTorch 与静态框架最大的区别是动态图而动态图在源码里的物理载体是每个 Tensor 上的 AutogradMeta 结构。简单说一个 Tensor 不仅保存数值和 shape还会带上一个指向 grad_fn 的指针grad_fn 代表“产生这个张量的运算”比如加法产生的 AddBackward0。TensorImpl ├── autograd_meta_ → AutogradMeta │ ├── requires_grad_ │ ├── grad_fn_ → 指向 Node如 AddBackward0 │ └── grad_ → 当前累计梯度这段是目录结构之外最值得背下来的抽象每一次运算只要输入张量里有 requires_gradTrue 的框架就会在输出的 Tensor 上挂一个新的 grad_fn。反向传播时从最终 loss 的 grad_fn 出发沿着 Node 之间的边往前遍历每经过一个节点调用它对输入求梯度的函数。这个图是在执行过程中实时构建出来的所以叫动态图。有了这张图后面精读 backward 的时候就不容易迷路。3. 跑通与调试构建一个能单步穿透到 C 的源码环境很多读者拿到源码后第一步就是打开 IDE 看代码但我强烈建议先把环境跑起来因为 PyTorch 源码里 Python 与 C 的边界很厚纯看代码容易漏掉启动顺序和绑定关系。用预编译的 wheel 包做调试往往只能停在 Python 层断点永远进不了 C。为了能单步穿透先花点时间从源码构建一份带调试信息的本地版本后面所有精读都会省力很多。3.1 环境选择什么时候该从源码构建如果只是想读 Python 层逻辑直接用 pip 装的 pyTorch 加 pdb 就够了。但如果你要改 C 侧代码或者想看到 Dispatcher 和 kernel 的真实调用栈就必须自己从源码构建。预编译 wheel 是 release 版本符号表不完整而且 C 文件路径被剥离gdb 里看到的是一堆二进制地址没法映射回源码行号。下面的命令是我常用的源码构建流程建议在无 GPU 的机器上也顺一遍因为 CPU 构建能验证整条工具链conda create -n torch-src python3.10 -y conda activate torch-src pip install cmake ninja ccache git clone https://github.com/pytorch/pytorch.git cd pytorch pip install -r requirements.txt DEBUG1 python setup.py develop先解释每个步骤的作用。conda 单独建环境是为了隔离编译器和 Python 版本避免污染其他项目cmake、ninja、ccache 是编译依赖ninja 并行构建比 make 快不少ccache 用于缓存中间编译产物二次改动 C 源码后增量编译能快一个数量级。DEBUG1 会开启调试符号和关闭部分优化代价是构建时间更长、运行更慢但对断点调试来说是必要的。最后用 python setup.py develop 而不是 pip install .develop 模式会把源码目录软链到 site-packages之后改 Python 文件不需要重装。提示首次完整构建在普通多核 CPU 机器上可能耗时 30 到 60 分钟如果只是想验证流程可以先用 CPU-only 的最小配置跑通再决定要不要加 CUDA 支持。构建期间不要切换 conda 环境GCC 和 CUDA 版本混用是常见的翻车来源。3.2 用 pdb 把第一次 backward 断住构建完成之后先不要急着进 C而是用一个最小训练脚本验证 Python 侧断点链路正常。把下面的代码保存成 train_min.pyimport torch def main(): x torch.randn(4, 4, requires_gradTrue) w torch.randn(4, 4, requires_gradTrue) y (x * w).sum() y.backward() # 想看到 autograd 引擎启动在这里打断点 print(w.grad) if __name__ __main__: main()用 pdb 启动并下断点python -m pdb train_min.py (Pdb) break train_min.py:8 (Pdb) continue (Pdb) bt断点命中后使用 bt 查看调用栈你会看到 y.backward() 一层层走进 torch/autograd/init.py 里的 backward 函数它校验 grad_tensors 和 retain_graph 参数后最终调用 _C._base_autograd_backward。这一行是从 Python 跨进 C 的分水岭再往下走Python 的 pdb 就看不到堆栈了因为执行已经交给了 C 引擎你只能在 gdb 里继续追踪。这里值得记一下参数语义retain_graph 为 True 时不释放中间节点的图结构默认情况下 backward 执行完会清掉计算图以省内存所以第二次调用同样的 backward 会报错提示 retain_graphcreate_graph 为 True 时这次 backward 本身也会构建图用于计算二阶导。理解这两个参数的意思是读懂后续反向传播代码的前提。3.3 需要深入 C 时用 gdb 看 kernel 与 autograd 引擎当 pdb 走到边界之后换成 gdb 继续看 C 侧。启动方式很简单gdb --args python train_min.py (gdb) break THPVariable_add (gdb) run (gdb) btTHPVariable_add 是 pybind11 为 Tensor.add 生成的 C 绑定函数打断点后每次张量加法都会停在这里bt 可看到 ATen dispatcher 和 kernel 的完整调用链。如果你用的是本文 3.1 节的 DEBUG 构建源代码行号会完整展示否则 break 会直接失败。在 gdb 里比较难观察的是 Python 对象所以实际调试时常用策略是 pdb 确认 Python 逻辑gdb 确认 C 调度两个工具配合而不是互相替代。从实践看大部分二次开发其实不需要走到 gdb 这一步python 层调度链看清楚就够用。但有一个场景绕不开 gdb你怀疑某次崩溃发生在异步 CUDA kernel 里CPU 堆栈已经返回错误却在后续同步时才爆出来。这种问题光靠 pdb 定位不了。4. 核心模块精读autograd 引擎与 nn.Module 的关键实现地图和工具都备好了进入正题。这一章挑四个最常被问到的核心点拆开讲顺序照着数据流走先看 Tensor 的 AutogradMeta 怎么构建图节点再看 backward 从 Python 到 C 的入口然后看 nn.Module 怎么组织参数最后看优化器怎么消费梯度。4.1 Tensor.grad_fn 与 AutogradMeta动态图的节点信息先看动态图的最小结构。PyTorch 里每个 Tensor 的核心对象是 TensorImpl而 autograd 相关的状态存在它指向的 AutogradMeta 里核心字段可以用下面的结构来理解struct TensorImpl { AutogradMeta* autograd_meta_ nullptr; // 其他字段sizes、strides、data_ptr 等 }; struct AutogradMeta { bool requires_grad_ false; bool is_view_ false; // 是否是某个张量的视图 std::shared_ptrNode grad_fn_; // 记录产生本张量的运算节点 Variable grad_; // 累计的梯度 };这段代码的精髓在 grad_fn_。当 x.requires_gradTrue 参与乘法运算时输出张量的 AutogradMeta 会被设置为指向一个 MulBackward0 节点这个节点保存了前向时的输入引用用于反向计算。反向传播会沿 grad_fn 的 next_edges 递归前行每到一个节点就调用它的 apply 方法计算对输入的梯度并把结果累加到前一个节点的 grad_ 里。关于 requires_grad_ 和 grad_fn_ 的关系有一个常见误区requires_grad 为 False 的张量不一定没有 grad_fn比如一个 requires_gradTrue 的输入经过 view 操作产生的输出requires_grad 也可以是 False但它仍指向一个 ViewBackward 节点。判断一个张量是否参与反向传播正确方式是检查 grad_fn 是否为空而不是看 requires_grad。4.2 backward 入口Python 封装与 C 引擎的分工torch/autograd/init.py 里的 backward 函数是整个反向传播的 Python 入口它做的事比想象中更薄。把关键逻辑浓缩出来大概是这样的# torch/autograd/__init__.py 里的函数做参数校验和默认值处理 def backward(tensors, grad_tensorsNone, retain_graphNone, create_graphFalse): # 传了 grad_tensors 时要检查形状是否匹配 if grad_tensors is None: grad_tensors [None] * len(tensors) if retain_graph is None: # 默认情况下动态图用完释放 retain_graph create_graph # 真正进入 C 引擎 _C._base_autograd_backward( tensors, grad_tensors, retain_graph, create_graph, allow_unreachableTrue, )参数说明grad_tensors 是每个根张量的初始梯度不传时引擎会为标量 loss 自动填一个 1.0 的梯度retain_graph 控制计算图是否在反向后保留create_graph 决定是否把反向过程本身也构造成图以便计算高阶导数。真正遍历图、调用每个 Node 的 apply 方法、累加梯度的逻辑全部在 C 引擎里Python 端只是把参数组装好再交给 _base_autograd_backward。这里暴露的边界感对二次开发很重要如果你想在反向传播每个算子执行前后插入日志最佳位置不是改这个 Python 函数而是给某个自定义 Node 写一个 apply 方法或者在 dispatcher 层挂 kernel 回调。改 Python 的 backward 入口只会影响入口改不到每个算子的执行路径。4.3 nn.Module 如何递归管理参数与子模块nn.Module 是模型搭建的基石它用两个有序字典分别管理参数和子模块。把构造函数和参数注册逻辑浓缩成下面的结构class Module: def __init__(self): self._parameters OrderedDict() # 本模块直接持有的参数 self._modules OrderedDict() # 子模块用于递归遍历 def register_parameter(self, name, param): if param is None: # 允许占位比如模型中可选分支的权重 self._parameters[name] None elif not isinstance(param, Parameter): # Parameter 是 Tensor 子类且 requires_grad 默认为 True raise TypeError(...) self._parameters[name] paramParameter 与普通 Tensor 唯一的实质区别是它创建时就把 requires_grad 设为 True这样注册进 Module 后所有.requires_grad的逻辑都能统一走参数路由。当你调用 module.parameters() 时框架会同时遍历 _parameters 和 _modules对子模块做递归最终返回一个扁平列表。这个递归过程在源码里落到 named_parameters 函数它用遍历加去重的方式保证同一个参数不会因为被多个模块引用而重复返回。给自定义层写代码时如果直接用 self.weight torch.randn(...) 而不是包成 Parameter权重不会进入优化器的参数列表训练时就不会更新这是新手最常见的坑之一。所有需要被优化器更新的量都必须显式创建为 torch.nn.Parameter。4.4 优化器 step 如何消费梯度优化器与 autograd 的衔接点在于反向传播把梯度累加到 parameter.grad而优化器 step 读取这个 grad 做原地更新。看 SGD 的核心循环就能理解整个优化器家族的模式def step(self): for group in self.param_groups: lr group[lr] weight_decay group[weight_decay] for p in group[params]: if p.grad is None: continue # 没有梯度的参数跳过 g p.grad.data if weight_decay ! 0: g g p.data * weight_decay # 用 data 原地更新不经过 autograd避免再建图 p.data.add_(-lr, g)这段代码有两个关键点。第一个是 p.grad 的形态梯度在反向传播中是被累加的所以每个 step 之前需要调用 optimizer.zero_grad() 把梯度清零否则前一个 batch 的梯度会和其他 batch 叠加。第二个关键点是为什么更新时要用 p.data 而不是直接操作 pp.data 是底层张量的裸数据对它做 add_ 可以避免 autograd 为这次更新再构建一个反向节点如果直接对 p 做运算计算图会膨胀还会破坏反向传播的语义。这也是在源码里反复出现 data 的原因。5. 常见问题与排查读 PyTorch 源码时容易翻车的几个细节源码调试和业务开发不同它的“玄学”集中在构建、符号、跨语言边界这几个地方。把最常碰到的五个问题整理成排查记录每一条都是踩过的坑现象、原因、解决三步说清楚。5.1 现象改了 Python 端 requires_grad 逻辑运行却不生效现象在 torch/_tensor.py 或者自定义模块里修改 requires_grad 的赋值逻辑重启进程后行为完全没变甚至直接报错。原因Tensor 的 requires_grad 真源在 C 侧的 AutogradMetaPython 端属性只是一个代理。你改的是代理C 侧原始状态没动。解决要么通过 torch.no_grad() 等官方上下文管理要么真正修改 C 源码后重新编译。自测代码里不要尝试直接赋值 requires_grad 来绕过防护逻辑。5.2 现象pdb 能进 Python 方法gdb 断点却停在函数入口就丢掉现象用 gdb 在 THPVariable_add 打断点break 能设置上但 run 之后要么没命中要么命中了没有源代码行号。原因预编译 wheel 包不包含调试源码映射符号也被 strip 过运行时代码路径与原仓库路径不一致。解决使用源码构建并开启 DEBUG1。构建完成后在 gdb 里用 directory 手动指定源码路径也可以救回一部分符号但最可靠的还是从一开始就用 DEBUG 构建。5.3 现象改了 torch/autograd/ 下的 Python 文件重启后还是老逻辑现象修改 torch/autograd/init.py 想加一行打印运行脚本却没有任何输出。原因常见有两个一是 develop 模式没有重新执行导致二进制模块缓存二是 Python 的pycache缓存了旧的 pyc。解决确认你是用 setup.py develop 安装改动后删除 site-packages 下对应目录的pycache然后重新启动一个新的 Python 进程不要在旧会话里 reimport。5.4 现象源码构建的版本拷贝到另一台机器后一跑就段错误现象在构建机上正常把整个 conda 环境或源码目录拷到别的机器上import torch 直接崩溃。原因源码构建的扩展模块包含当前机器的编译器和 CPU 指令集相关信息比如 GCC 版本对应 ABI 不一致。解决云开发环境的常见做法是每台同规格机器独立重新构建不要拷贝二进制借助 ccache 可加速重复构建。排查时优先用 gdb 看到崩溃点是不是在 glibc 或者编译器 runtime 里。5.5 现象改 native_functions.yaml 新增算子后编译报 schema 冲突现象加了一行算子声明编译到 codegen 阶段报 duplicate schema 或者缺少签名。原因native_functions.yaml 是代码生成器的唯一事实来源C kernel 的签名必须和 yaml 里的 schema 完全一致参数顺序、optional 标记不能错。解决别手写绑定代码直接在 yaml 里声明后重新构建codegen 会自动生成 C 和 Python 两侧的模板代码。写 kernel 前先复制一个同类型算子的 yaml 条目和 kernel 定义改参数名而不是从零造。6. 二次开发实战注册一个自定义算子并验证梯度正确性最后的改造走一条最常见的路径注册一个自定义算子并确保前向、反向都能正确工作。先做 Python 级的实现成本最低也方便和其他模块对接。6.1 自定义 autograd.Functionimport torch class MyLinearFunction(torch.autograd.Function): staticmethod def forward(ctx, x, w, b): # save_for_backward 保存的 Tensor 在反向时仍然可读 ctx.save_for_backward(x, w) return x.mm(w.t()) b staticmethod def backward(ctx, grad_output): x, w ctx.saved_tensors grad_x grad_output.mm(w) grad_w grad_output.t().mm(x) grad_b grad_output.sum(0) return grad_x, grad_w, grad_bbackward 的返回值必须和 forward 的输入参数一一对应x、w、b 三个输入对应返回三个梯度。save_for_backward 只能保存 Tensor不能保存 Python 标量b 没有参与矩阵运算直接对 grad_output 求和就能得到它的梯度。验证梯度时需要走 gradcheckfrom torch.autograd import gradcheck x torch.randn(4, 3, dtypetorch.double, requires_gradTrue) w torch.randn(5, 3, dtypetorch.double, requires_gradTrue) b torch.randn(5, dtypetorch.double, requires_gradTrue) inputs (x, w, b) check gradcheck(MyLinearFunction.apply, inputs, eps1e-6, atol1e-4) print(gradcheck passed:, check)gradcheck 用数值差分和你的反向实现做对比eps 是差分步长atol 是容差。它是验证反向传播正确性的金标准任何自定义算子提交前都应该跑一次。从那以后我每次改完源码都强制走一遍 gradcheck 加一个最小训练脚本确认参数确实更新才继续往业务层叠。希望帮到你。本文还有配套的精品资源点击获取