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

PyTorch 深度学习框架与 GPU 加速生态：从入门到理解整个技术栈

news 2026/7/3 12:18:33

很多初学者在安装 PyTorch 时，都会遇到诸如CUDA、ROCm、MPS、cuDNN、显卡驱动、Python 版本等各种专业术语。它们之间到底是什么关系？为什么有的人安装了 CUDA，有的人却不用安装？为什么同样都是 GPU，不同品牌却需要不同版本的 PyTorch？

本文将从整体架构出发，系统梳理 PyTorch 的 GPU 加速生态，帮助你建立完整的知识体系。

一、PyTorch：深度学习框架的核心

PyTorch是目前最流行的开源深度学习框架之一，由 Facebook AI Research（FAIR，现 Meta AI）开发，目前由 PyTorch Foundation 维护。

很多资料会说：

“PyTorch 基于 Torch 开发。”

这种说法并不准确。

更准确地说：

PyTorch 借鉴了早期 Lua Torch（Torch7）的设计思想，但底层已经完全重构，核心采用 C++（LibTorch/ATen）实现，而 Python 只是最常用的接口。

PyTorch 提供了：

Tensor（张量）计算
自动求导（Autograd）
神经网络模块（torch.nn）
数据加载（DataLoader）
模型训练与推理
分布式训练
GPU 加速

其中，Tensor（张量）是 PyTorch 的核心数据结构。

它与 NumPy ndarray 非常相似，但最大的区别在于：

可以进行自动求导
可以直接运行在 GPU 上
支持 CPU、CUDA、ROCm、MPS 等多种计算后端

例如：

importtorch x=torch.randn(3,3)print(x.device)# cpux=x.to("cuda")print(x.device)# cuda:0

这里真正发生的是：

Python 代码 → PyTorch → GPU 后端 → GPU 硬件。

二、PyTorch 内部架构

很多人认为 PyTorch 的执行流程非常简单：

Python ↓ PyTorch ↓ CUDA ↓ GPU

实际上，PyTorch 内部还有多个重要模块：

Python API │ ┌──────────────┴──────────────┐ │ │ torch.nn torch.optim │ Autograd（自动求导） │ Dispatcher（统一调度器） │ ├──────── CPU Backend ├──────── CUDA Backend ├──────── ROCm Backend └──────── MPS Backend │ 各种计算库 │ GPU / CPU

其中：

1）ATen

ATen 是 PyTorch 的底层张量计算库。

几乎所有 Tensor 运算都会最终调用 ATen。

例如：

a+b a @ b torch.matmul()torch.conv2d()

最终都会转换成底层计算接口。

2）Autograd

Autograd 是 PyTorch 的自动求导系统。

例如：

loss.backward()

梯度就是由 Autograd 自动计算。

这也是 PyTorch 能够训练神经网络的重要基础。

3）Dispatcher

Dispatcher 可以理解成：

PyTorch 的"交通指挥中心"。

例如：

x+y

Dispatcher 会自动判断：

如果数据在 CPU：

CPU Backend

如果数据在 NVIDIA GPU：

CUDA Backend

如果数据在 AMD GPU：

ROCm Backend

如果数据在 Apple Silicon：

MPS Backend

因此，大多数情况下，开发者无需修改算法代码，只需要切换设备即可。

三、四种计算后端（Backend）

Backend（后端）负责将 PyTorch 的运算转换成对应硬件能够执行的指令。

后端	对应硬件	主要生态
CPU	所有 CPU	oneDNN、MKL、OpenMP
CUDA	NVIDIA GPU	CUDA
ROCm	AMD GPU	HIP、rocBLAS、MIOpen
MPS	Apple Silicon	Metal Performance Shaders

可以理解为：

PyTorch │ ├── CPU ├── CUDA ├── ROCm └── MPS

四、CUDA：NVIDIA GPU 的计算平台

CUDA（Compute Unified Device Architecture）是 NVIDIA 推出的 GPU 并行计算平台和编程模型，也是目前最成熟、生态最完善的 GPU 计算方案。

需要注意的是：

CUDA 并不仅仅是一个软件，而是一整套 GPU 软件生态。

它主要包括：

CUDA Toolkit │ ├── CUDA Runtime ├── CUDA Driver API ├── cuBLAS ├── cuDNN ├── NCCL ├── nvcc └── 其它开发工具

其中：

CUDA Runtime：程序运行所需的运行时库
CUDA Driver API：负责与显卡驱动通信
cuBLAS：高性能矩阵运算库
cuDNN：深度学习加速库
NCCL：多 GPU 通信库
nvcc：CUDA 编译器

因此，PyTorch 并不是直接与 GPU 通信，而是通过 CUDA Runtime、cuBLAS、cuDNN、NCCL 等组件完成 GPU 加速。

五、CUDA Runtime 与 CUDA Toolkit 的区别

这是初学者最容易混淆的地方。

很多人认为：

使用 PyTorch GPU 必须安装 CUDA Toolkit。

事实上并非如此。

需要区分三个概念：

组件	是否必须	作用
NVIDIA Driver	✅ 必须	驱动 GPU
CUDA Runtime	✅ 必须（通常随 PyTorch 提供）	GPU 运行环境
CUDA Toolkit	❌ 大多数用户无需安装	CUDA 开发工具

对于绝大多数用户：

pipinstalltorch torchvision

安装的实际上是：

torch ├── CUDA Runtime ├── cuBLAS ├── cuDNN ├── NCCL

因此，即使系统没有安装 CUDA Toolkit，PyTorch 依然可以正常调用 GPU。

只有在以下场景才需要安装 CUDA Toolkit：

编写 CUDA 程序
使用 nvcc 编译 CUDA 代码
编译 PyTorch CUDA Extension
开发 CUDA 算法

六、cuDNN、cuBLAS 与 NCCL

很多人误以为 CUDA 完成了所有 GPU 计算。

实际上，CUDA 更像一个平台，真正承担大量工作的，是多个专用计算库。

1）cuBLAS

负责：

矩阵乘法
向量计算
GEMM 运算

Transformer、MLP、Linear 等大量运算都依赖 cuBLAS。

2）cuDNN

负责深度学习相关算子，例如：

卷积（Conv）
池化（Pooling）
BatchNorm
RNN
LSTM
Transformer 部分算子

训练 CNN 或 Transformer 时，大量计算都会调用 cuDNN。

3）NCCL

NCCL（NVIDIA Collective Communications Library）负责多 GPU 通信。

例如：

DistributedDataParallel（DDP）
FSDP
多 GPU 训练
多机训练

都会大量使用 NCCL 完成：

AllReduce
Broadcast
Gather
ReduceScatter

因此，对于大模型训练来说，NCCL 是不可或缺的重要组件。

七、ROCm：AMD GPU 的官方计算平台

ROCm（Radeon Open Compute）是 AMD 推出的 GPU 计算平台。

可以把它理解成：

CUDA 在 AMD 世界中的对应方案。

其生态包括：

ROCm │ ├── HIP ├── rocBLAS ├── MIOpen ├── RCCL

分别对应 CUDA 世界中的：

CUDA	ROCm
CUDA Runtime	HIP
cuBLAS	rocBLAS
cuDNN	MIOpen
NCCL	RCCL

目前，ROCm 已支持越来越多的 Radeon 与 Instinct 系列 GPU，但整体生态成熟度、第三方软件支持以及模型适配速度仍与 CUDA 存在一定差距。

八、MPS：Apple Silicon 的 GPU 后端

对于搭载 Apple Silicon（M 系列芯片）的 Mac，PyTorch 使用MPS（Metal Performance Shaders）作为 GPU 后端。

可以通过以下方式检测是否可用：

torch.backends.mps.is_available()

使用方式：

device=torch.device("mps")

MPS 已能够支持大多数常见模型训练与推理，但相比 CUDA 仍存在一些限制：

部分算子尚未实现
某些操作可能自动回退到 CPU
多 GPU 支持有限
生态成熟度仍低于 CUDA

因此，在苹果平台上，MPS 非常适合开发、学习和中小规模模型实验，而在大规模训练场景中，CUDA 仍是主流选择。

九、CPU 后端并非没有优化

没有 GPU 时，PyTorch 依然可以充分利用现代 CPU 的计算能力。

不同平台通常会调用不同的优化库，例如：

平台	常见优化库
Intel CPU	oneDNN、MKL
Linux	OpenMP
Apple CPU	Accelerate Framework

因此，即使在 CPU 环境下，PyTorch 也能获得较好的计算性能。

十、显卡驱动：整个 GPU 生态的基础

显卡驱动是操作系统与 GPU 硬件之间的底层通信接口。

它主要负责：

识别 GPU
管理显存
调度 GPU
为 CUDA Runtime 提供底层接口

对于 NVIDIA GPU 而言：

驱动版本决定了系统能够支持的最高 CUDA Runtime 版本。

需要注意的是：

驱动版本号（如 551.xx、560.xx）与 CUDA 版本号（如 12.4、12.6）属于两套不同的版本体系，不能直接对应。

例如：

CUDA 12.4 需要 Driver ≥ 550.xx

不同 CUDA 版本都有对应的最低驱动要求，应以 NVIDIA 官方兼容矩阵为准。

十一、理解 nvidia-smi 与 nvcc 的区别

这是环境配置中最容易产生误解的问题之一。

假设执行：

nvidia-smi

输出：

CUDA Version: 12.6

很多人会误认为：

我安装了 CUDA 12.6。

实际上，这表示的是：

当前 NVIDIA 驱动最高支持 CUDA 12.6 Runtime。

如果想查看系统是否安装了 CUDA Toolkit，应执行：

nvcc--version

因此：

命令	查看内容
nvidia-smi	驱动版本、GPU 信息、驱动支持的最高 CUDA Runtime
nvcc --version	已安装的 CUDA Toolkit 版本

两者代表的含义完全不同。

十二、Python 与 PyTorch 的版本兼容性

PyTorch 是 Python 库，因此必须与 Python 版本兼容。

通常：

不同版本的 PyTorch 支持不同的 Python 版本范围。
新版本 PyTorch 一般要求 Python 3.10 或更高版本（具体以官方发布说明为准）。
建议使用官方支持范围内的 Python 版本，并优先参考安装页面的推荐组合。

为了避免不同项目之间发生依赖冲突，建议使用：

venv
Conda
Miniforge

等虚拟环境进行隔离。

十三、PyTorch 的设备（Device）抽象

PyTorch 提供了统一的设备接口，使同一套代码可以运行在不同硬件上。

例如：

device=torch.device("cpu")device=torch.device("cuda")device=torch.device("cuda:0")device=torch.device("mps")

模型：

model.to(device)

数据：

data.to(device)

因此，大多数情况下，只需修改设备即可在 CPU、NVIDIA GPU、AMD GPU 或 Apple Silicon GPU 之间切换，而无需改动模型逻辑。

十四、PyTorch 安装包命名说明

安装 PyTorch 时，经常会看到如下版本：

torch==2.x.x+cpu torch==2.x.x+cu124 torch==2.x.x+rocm6.x

其中：

后缀	含义
+cpu	CPU 版本
+cu124	内置 CUDA 12.4 Runtime 的版本
+rocm6.x	ROCm 版本

这些后缀表示PyTorch 所携带的运行时后端，而不是要求系统必须安装对应版本的 CUDA Toolkit。

十五、推荐的安装方式

由于 PyTorch、Python、CUDA、ROCm 等版本之间存在兼容关系，最推荐的方式是使用 PyTorch 官方安装页面自动生成安装命令。

用户只需选择：

操作系统
包管理器（pip 或 Conda）
Python 环境
计算平台（CPU、CUDA、ROCm）

即可获得对应的安装命令，避免因版本选择错误导致环境配置失败。

十六、完整的软件栈与兼容关系

从底层硬件到上层应用，一个完整的 PyTorch GPU 计算栈如下：

Python │ PyTorch API │ ┌──────── Dispatcher ────────┐ │ │ │ CPU CUDA ROCm / MPS │ │ │ oneDNN/MKL cuBLAS、cuDNN、NCCL HIP、MIOpen、Metal └─────────────┴──────────────┘ │ 显卡驱动（Driver） │ GPU（NVIDIA / AMD / Apple）

整个计算链路可以概括为：