学 PyTorch，第一个门槛不是模型。是环境。环境装错，代码还没开始跑，报错已经堆满屏。

这一章只讲一件事：把 PyTorch 环境一次性讲清楚。不是背安装命令，而是看懂安装背后的逻辑。

因为命令会随着版本变化。逻辑不会。你只要能判断 CPU、CUDA、ROCm、MPS 该选谁，后面的训练、调试、部署才不会乱。

一、PyTorch 安装，真正安装的是什么

很多人以为安装 PyTorch 就是：pip install torch。

不对。真正安装的是一个“带后端能力的 torch 包”。它可能是 CPU 构建，也可能是 CUDA 构建，也可能是 ROCm 构建。Mac 上的 MPS 则走 Apple Metal 后端。

所以，安装 PyTorch 不是问“装哪个版本”。而是先问：你的计算要交给谁执行？CPU？NVIDIA GPU？AMD GPU？还是 Apple Silicon 的 Metal GPU？

二、CPU 版：最稳，最慢，最适合起步

CPU 版适合三类场景：刚入门、写教程、先跑通代码。

它的优点是简单。不依赖显卡，不依赖驱动，不挑机器。只要 Python 环境正常，基本能跑。

它的缺点也明显：慢。小模型能忍，大模型直接劝退。

# CPU 环境验证
import torch
x = torch.rand(5, 3)
print(x)
print(x.device) # cpu

记住一句话：CPU 版适合把路打通，不适合把大模型训起来。

三、CUDA 版：NVIDIA 显卡的主流路线

只要你是 NVIDIA 显卡，深度学习训练优先考虑 CUDA。

CUDA 版 PyTorch 的关键不在“有没有显卡”，而在四件事是否匹配：PyTorch 构建、CUDA 运行时、NVIDIA 驱动、GPU 架构。

最常见的坑是：机器明明有 NVIDIA 显卡，torch.cuda.is_available() 却返回 False。多数时候，不是模型错了，是环境链路断了。

# CUDA 环境验证
import torch
print(torch.__version__)
print(torch.version.cuda)
print(torch.cuda.is_available())
if torch.cuda.is_available():
print(torch.cuda.get_device_name(0))
x = torch.rand(5, 3, device="cuda")
print(x.device)

官方安装页的做法是：在选择器里按操作系统、包管理器、Python 和 Compute Platform 生成安装命令。不要复制旧文章里的固定命令。

四、ROCm 版：AMD 显卡路线，先看兼容

AMD 显卡不是不能做深度学习。它走 ROCm。

但 ROCm 对系统、显卡型号、驱动版本的要求更严格。尤其是生产环境，建议优先考虑 Docker。AMD 官方文档也把预装 PyTorch、ROCm 和依赖的 Docker 镜像作为推荐方式。

# ROCm 环境下，Python API 仍然常用 torch.cuda 语义
import torch
print(torch.version.hip) # ROCm 构建通常会有值
print(torch.cuda.is_available()) # ROCm 下也沿用 cuda API 语义

注意：ROCm 版里很多 Python 层 API 仍然叫 torch.cuda。这不是写错，而是 PyTorch 为了减少上层代码差异，沿用了相近语义。

五、MPS 版：Apple Silicon 的本地加速

如果你是 M1、M2、M3、M4 系列 Mac，可以用 MPS 后端。MPS 基于 Apple Metal，适合本地训练、调试、小规模实验。

但不要把 MPS 当成 CUDA 的完全替代。它能加速很多模型，但算子覆盖、性能表现、内存行为和 CUDA 不完全一样。

# MPS 环境验证
import torch
if torch.backends.mps.is_available():
device = torch.device("mps")
else:
device = torch.device("cpu")
x = torch.ones(5, device=device)
print(x.device)

MPS 里要分清两个问题：这个 torch 包有没有编进 MPS？当前机器和系统能不能真正启用 MPS？前者看 is_built，后者看 is_available。

六、源码级看一眼：is_available() 不是魔法

环境排查不要只看表面。我们直接看源码思路。

CUDA 的 is_available 大体分两层：第一层看这个 PyTorch 是否编译了 CUDA 支持；第二层调用底层 C++ 接口检查设备数量。

# torch/cuda/__init__.py 核心逻辑简化
def _is_compiled():
return hasattr(torch._C, "_cuda_getDeviceCount")
def is_available():
if not _is_compiled():
return False
return torch._C._cuda_getDeviceCount() > 0

这就解释了为什么 False 不等于“你没有显卡”。它可能只是当前 torch 包根本不是 CUDA 构建。

MPS 的逻辑更直接：is_built 看构建能力，is_available 看运行时是否可用。

# torch/backends/mps/__init__.py 核心逻辑简化
def is_built():
return torch._C._has_mps
def is_available():
return torch._C._mps_is_available()

七、为什么 CUDA 报错经常“慢半拍”

GPU 操作不是普通 Python 函数。PyTorch 调用 GPU 时，很多操作会先进入队列，然后由 GPU 异步执行。

这就会出现一个现象：真正出错的地方，可能不是报错显示的那一行。因为错误是在 GPU 后面执行时才暴露出来。

调试时可以临时设置 CUDA_LAUNCH_BLOCKING=1，让 GPU 操作同步执行。这样报错位置更接近真实原因。

# Linux / macOS
CUDA_LAUNCH_BLOCKING=1 python train.py
# Windows PowerShell
$env:CUDA_LAUNCH_BLOCKING=1
python train.py

八、显存不是你以为的那样释放

很多人看到 nvidia-smi 显示显存没降，就以为 PyTorch 没释放。

实际上，PyTorch 有 CUDA caching allocator。它会缓存释放后的显存块，让下一次分配更快。

所以，显存排查要分清两个概念：Tensor 真正占用的显存，和 PyTorch 缓存管理器保留的显存。

# 显存观察
import torch
print(torch.cuda.memory_allocated() / 1024**2, "MB")
print(torch.cuda.memory_reserved() / 1024**2, "MB")
# 释放未使用缓存，不会释放仍被 Tensor 占用的显存
torch.cuda.empty_cache()

memory_allocated 看 Tensor 正在占多少；memory_reserved 看缓存分配器总共管着多少。不要只盯 nvidia-smi。

九、开发环境的正确姿势

不要把所有项目都装进一个 Python。

一个项目一个虚拟环境。一个后端一套依赖。CPU、CUDA、ROCm 不要混在一个环境里反复覆盖。

真正工程化的 PyTorch 项目，必须能复现环境。否则同一份代码，今天能跑，明天就炸。

十、常见问题，一张图排查

环境问题不要凭感觉修。按症状定位，速度最快。

十一、总结

• 第一，先选硬件后端，再复制安装命令。

• 第二，CPU 版最稳，CUDA 版最主流，ROCm 要看兼容，MPS 适合 Mac 本地调试。

• 第三，安装成功不等于 GPU 可用，必须用代码验证。

• 第四，is_available 返回 False 时，优先排查 wheel、驱动、设备可见性。

• 第五，显存和 CUDA 报错有异步与缓存机制，不能只看表面。

下一章开始进入 PyTorch 的真正核心：Tensor。Tensor 是所有模型、梯度、参数和计算图的起点。