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

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

    • 伪代码

      python 复制代码
      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

相关推荐
秋千码途17 分钟前
小架构step系列26:Spring提供的validator
java·spring·架构
大数据张老师19 分钟前
用 AI 做数据分析:从“数字”里挖“规律”
大数据·人工智能
音视频牛哥1 小时前
如何打造毫秒级响应的RTSP播放器:架构拆解与实战优化指南
人工智能·机器人·音视频开发
张较瘦_1 小时前
[论文阅读] 人工智能 + 软件工程 | NoCode-bench:评估LLM无代码功能添加能力的新基准
论文阅读·人工智能·软件工程
go54631584651 小时前
Python点阵字生成与优化:从基础实现到高级渲染技术
开发语言·人工智能·python·深度学习·分类·数据挖掘
Coovally AI模型快速验证1 小时前
避开算力坑!无人机桥梁检测场景下YOLO模型选型指南
人工智能·深度学习·yolo·计算机视觉·目标跟踪·无人机
西陵2 小时前
Nx带来极致的前端开发体验——借助playground开发提效
前端·javascript·架构
巫婆理发2222 小时前
神经网络(第二课第一周)
人工智能·深度学习·神经网络
Edingbrugh.南空2 小时前
Aerospike架构深度解析:打造web级分布式应用的理想数据库
数据库·架构
欧阳小猜2 小时前
OpenCV-图像预处理➁【图像插值方法、边缘填充策略、图像矫正、掩膜应用、水印添加,图像的噪点消除】
人工智能·opencv·计算机视觉