环境安装与开发姿势：CPU、CUDA、ROCm、MPS 怎么选

学 PyTorch第一个门槛不是模型。是环境。环境装错代码还没开始跑报错已经堆满屏。这一章只讲一件事把 PyTorch 环境一次性讲清楚。不是背安装命令而是看懂安装背后的逻辑。因为命令会随着版本变化。逻辑不会。你只要能判断 CPU、CUDA、ROCm、MPS 该选谁后面的训练、调试、部署才不会乱。一、PyTorch 安装真正安装的是什么很多人以为安装 PyTorch 就是pip install torch。不对。真正安装的是一个“带后端能力的 torch 包”。它可能是 CPU 构建也可能是 CUDA 构建也可能是 ROCm 构建。Mac 上的 MPS 则走 Apple Metal 后端。所以安装 PyTorch 不是问“装哪个版本”。而是先问你的计算要交给谁执行CPUNVIDIA GPUAMD GPU还是 Apple Silicon 的 Metal GPU二、CPU 版最稳最慢最适合起步CPU 版适合三类场景刚入门、写教程、先跑通代码。它的优点是简单。不依赖显卡不依赖驱动不挑机器。只要 Python 环境正常基本能跑。它的缺点也明显慢。小模型能忍大模型直接劝退。# CPU 环境验证import torchx 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 torchprint(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, devicecuda)print(x.device)官方安装页的做法是在选择器里按操作系统、包管理器、Python 和 Compute Platform 生成安装命令。不要复制旧文章里的固定命令。四、ROCm 版AMD 显卡路线先看兼容AMD 显卡不是不能做深度学习。它走 ROCm。但 ROCm 对系统、显卡型号、驱动版本的要求更严格。尤其是生产环境建议优先考虑 Docker。AMD 官方文档也把预装 PyTorch、ROCm 和依赖的 Docker 镜像作为推荐方式。# ROCm 环境下Python API 仍然常用 torch.cuda 语义import torchprint(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 torchif torch.backends.mps.is_available():device torch.device(mps)else:device torch.device(cpu)x torch.ones(5, devicedevice)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 Falsereturn 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_mpsdef is_available():return torch._C._mps_is_available()七、为什么 CUDA 报错经常“慢半拍”GPU 操作不是普通 Python 函数。PyTorch 调用 GPU 时很多操作会先进入队列然后由 GPU 异步执行。这就会出现一个现象真正出错的地方可能不是报错显示的那一行。因为错误是在 GPU 后面执行时才暴露出来。调试时可以临时设置 CUDA_LAUNCH_BLOCKING1让 GPU 操作同步执行。这样报错位置更接近真实原因。# Linux / macOSCUDA_LAUNCH_BLOCKING1 python train.py# Windows PowerShell$env:CUDA_LAUNCH_BLOCKING1python train.py八、显存不是你以为的那样释放很多人看到 nvidia-smi 显示显存没降就以为 PyTorch 没释放。实际上PyTorch 有 CUDA caching allocator。它会缓存释放后的显存块让下一次分配更快。所以显存排查要分清两个概念Tensor 真正占用的显存和 PyTorch 缓存管理器保留的显存。# 显存观察import torchprint(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 是所有模型、梯度、参数和计算图的起点。内容来源环境安装与开发姿势CPU、CUDA、ROCm、MPS 怎么选功能变化与行业影响解析_热闻岛