pytorch底层原理学习--PyTorch 架构梳理

PyTorch 架构梳理

PyTorch 完整架构流程图

硬件层

后端层

C++ 部署层

核心引擎 (libtorch C++)

绑定层

Python 层

加载

调用

训练模式

编译模式

推理模式

生成

CPU

GPU

CPU Kernels

CUDA Kernels

C++ 代码

torch::jit::load('model.pt')

module.forward(inputs)

libtorch C++ Frontend

训练路径

Autograd Engine

动态图记录

梯度计算

JIT 编译

语法解析

生成 IR

图优化

序列化 (.pt 文件)

JIT 执行

加载 IR

设备优化

无梯度执行

ATen (张量库)

Operator Dispatch

pybind11 Bindings

Python 代码

Python Frontend (torch, torch.nn)

Eager 操作

动态图构建

JIT 接口 (torch.jit.script/trace)

model.pt

Hardware

Core Engine

Binding Layer

Python Layer

Via

Via

CPU Instructions

CUDA Instructions
GPU Execution

C++ Frontend
libtorch

Autograd Engine
Dynamic Graph Construction

ATen
Tensor Library

Operator Dispatch

CPU Kernels
MKL/OpenMP/BLAS

CUDA Kernels
cuDNN/cuBLAS

TorchScript IR
Intermediate Representation

Graph Optimization
Fusion/Dead Code Elim

Graph Execution

pybind11 Bindings

Python Frontend
torch.nn, torch.optim, torch.Tensor

Python Code
Model/Training Scripts

Eager Execution
Immediate Operation Execution

Graph Execution
JIT/TorchScript

关键组件详解

1. Python Code

   • 作用：用户编写的模型定义、训练脚本
   • 示例：model = nn.Linear(10, 2); output = model(input)
   • 特点：高级API，易用性强

2. Python Frontend

   • 组成：torch, torch.nn, torch.optim等模块
   • 功能：提供神经网络层、优化器、张量操作等高级接口
   • 关键类：Tensor, Module, Optimizer

3. Eager Execution

   • 机制：命令式编程模式，操作立即执行
   • 优点：调试方便，动态图灵活性高
   • 示例：x = torch.tensor([1.0]); y = x * 2 (立即计算)

4. JIT/TorchScript

   • 作用：将Python模型编译为优化后的静态图
   • 流程：torch.jit.script(model) → 生成IR → 优化
   • 优势：部署友好，性能优化空间大

5. pybind11 Bindings

   • 功能：Python与C++间的双向绑定层
   • 实现：自动生成包装代码，实现无缝调用
   • 效率：接近原生C++性能的跨语言调用

6. Libtorch

   组件	功能
   TorchScript 支持	加载/执行 Python 导出的模型（.pt 文件）
   ATen 张量库	核心张量操作（CPU/CUDA）
   神经网络APIC++ Frontend	torch::nn 命名空间下的层实现
   自动求导引擎Autograd	C++ 环境下的 autograd 支持
   多后端支持C++ Extensions	CPU/CUDA/ROCM 硬件加速

   • 定位：PyTorch的C++核心库
   • 功能：提供与Python API对应的C++接口
   • 使用场景：高性能推理、嵌入式部署

7. Autograd Engine

   • 核心功能：动态构建计算图并管理梯度计算

8. ATen (A Tensor Library)

   • 角色：PyTorch的核心张量库
   • 特性：
     • 500+张量操作
     • 统一CPU/CUDA接口
     • 自动微分支持
   • 路径：aten/src/ATen/native/ (算子实现)

9. Intermediate Representation (IR)

   • 作用：TorchScript的中间表示形式
   • 结构：基于图的表示，包含节点(Node)、边(Edge)
   • 优化：常量折叠、算子融合等

10. Operator Dispatch

    • 机制：根据设备类型分发算子

    • 伪代码：

      def add(tensor):
          if tensor.device == 'cuda':
              return cuda_add_kernel(tensor)
          else:
              return cpu_add_kernel(tensor)
      

11. CPU Kernels

    • 优化技术：
      • SIMD指令集 (AVX2/AVX512)
      • 多线程并行 (OpenMP)
      • 数学加速库 (MKL, oneDNN)
    • 典型操作：矩阵乘法、卷积等

12. CUDA Kernels

    • 架构：

      CUDA Kernel

      线程块

      32线程束

      单个线程

    • 加速库：cuDNN (深度学习), cuBLAS (线性代数)

    • 异步执行：通过CUDA流实现计算/传输并行

完整执行流程示例

Python Code Python Frontend pybind11 C++ Frontend Autograd Engine ATen CUDA Kernel model(input) 张量操作请求 调用libtorch 记录前向图 执行张量计算 分发到CUDA内核 结果返回 存储梯度函数 返回输出张量 C++ → Python 返回结果 获得预测结果 Python Code Python Frontend pybind11 C++ Frontend Autograd Engine ATen CUDA Kernel

参考资料：

《deep learning with pytorch》15.3 与PyTorch JIT 编译器交互，第一版中文版

PyTorch

PyTorch Architecture | harleyszhang/llm_note | DeepWiki

(PyTorch源码分析（1）- 整体预览 - 知乎

《PyTorch: An Imperative Style, High-Performance Deep Learning Library》
https://pytorch.org/blog/a-tour-of-pytorch-internals-1/
https://pytorch.org/blog/a-tour-of-pytorch-internals-2/
https://zhuanlan.zhihu.com/p/598044604
https://zhuanlan.zhihu.com/p/708375873
https://zhuanlan.zhihu.com/p/338256656
https://blog.ezyang.com/2019/05/pytorch-internals/
https://deepwiki.com/harleyszhang/llm_note/5.2-pytorch-architecture?utm_source=chatgpt.com
https://medium.com/@hxu296/a-trip-to-kernels-understanding-pytorchs-internal-architecture-fc955aafd54c
https://zhuanlan.zhihu.com/p/609288586
https://mlfrontiers.substack.com/p/understanding-ml-compilers-the-journey
https://se.ewi.tudelft.nl/desosa2019/chapters/pytorch/#fnref:3
https://medium.com/@hxu296/a-trip-to-kernels-understanding-pytorchs-internal-architecture-fc955aafd54c
https://blog.christianperone.com/2018/03/pytorch-internal-architecture-tour/
https://docs.pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training