PyTorch 安装与部署

二、安装与部署

PyTorch 安装最容易踩的坑是「装上了但不能用 GPU」。根因通常是三件事没对齐：显卡驱动、wheel 里打包的 CUDA 运行时、以及 Python 版本。安装前先跑一条命令看清现状。

1. 确认显卡与驱动。执行 nvidia-smi，右上角会显示驱动版本与驱动支持的最高 CUDA 版本。这一步只需要驱动，不需要单独装 CUDA Toolkit——PyTorch 的 pip wheel 里自带所需 CUDA 运行时。
2. 决定安装方式。要复现性就用 conda/mamba 建独立环境；要轻量就用 venv + pip；要隔离得最彻底就用容器（官方 NGC 镜像里驱动由宿主机提供，容器内自带完整工具链）。
3. 按官方选择器拿命令。打开 pytorch.org 的安装选择器，选系统、包管理器、计算平台，复制它生成的命令。不要凭记忆写版本号——CUDA 版本组合每季度都在变，以选择器输出为准。
4. 创建环境并安装。下面是 pip 路径的骨架命令，CUDA 版本号请用选择器结果替换。
5. 验证。跑一段最小检查，确认版本、CUDA 可用性、以及张量能在 GPU 上做一次乘法。

# 1) 看驱动与驱动支持的最高 CUDA 版本
nvidia-smi

# 2) 建环境（二选一）
python -m venv .venv && source .venv/bin/activate      # Linux / macOS
conda create -n torch python=3.11 -y && conda activate torch

# 3) 装 PyTorch —— 下面这行的 --index-url 必须用官方选择器给出的实际值
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/<CUDA_TAG>

# 4) 验证
python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available(), torch.cuda.device_count())"

# 5) GPU 上真跑一次
python - <<'PY'
import torch
d = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(4096, 4096, device=d)
y = x @ x
print(d, y.shape, y.dtype)
PY

依赖与环境要求

• 操作系统：Linux 是主战场，Windows 与 macOS 官方支持。macOS 无 CUDA，只能用 MPS 后端（Apple 芯片的 GPU 加速）。
• Python 版本：跟随官方支持矩阵，旧版本 Python 会被新 wheel 抛弃，建议用仍处于官方支持周期内的版本。
• 驱动：驱动过旧会报「CUDA driver version is insufficient」；驱动足够新则向下兼容旧 CUDA 运行时。
• 磁盘：GPU 版 wheel 体积很大（含 CUDA 库），预留充足磁盘，pip 缓存也会占空间。
• 内存：DataLoader 的 num_workers 开大后，每个 worker 会复制一份进程内存，主机内存不足会直接被 OOM Killer 干掉。

常见报错与排查