Ubuntu PyTorch支持哪些GPU型号

Ubuntu 上 PyTorch 的 GPU 支持概览

在 Ubuntu 上，PyTorch 的 GPU 加速主要走两条路线：一是 NVIDIA CUDA 生态，二是 AMD ROCm 生态。是否能用取决于三件事：显卡架构与计算能力（NVIDIA）、PyTorch 预编译包所带的 CUDA/ROCm 版本、以及系统驱动版本是否匹配。

NVIDIA 显卡支持

• 已验证的主流架构与代表型号如下（不同 PyTorch/CUDA 版本覆盖略有差异，建议按表中“建议 PyTorch/CUDA”选配）：

• 实用提示
  • 运行 nvidia-smi 查看驱动与“支持的 CUDA 运行时版本”，再选择与之匹配的 PyTorch CUDA 包（无需单独安装完整 CUDA Toolkit，使用官方预编译包即可）。
  • 老卡（如 Kepler 及更早）已不在现代 PyTorch 支持范围。
  • 典型验证：import torch; print(torch.cuda.is_available()) 应为 True。

AMD 显卡支持

• ROCm 路线（Linux/Ubuntu 为主）：
  • 数据中心/加速卡：Instinct MI300/MI350 系列已在 ROCm 7.0 获得支持，适合大规模训练/推理。
  • 消费级显卡：Radeon RX 7900 XTX / RX 7900 XT / Radeon Pro W7900 等 RDNA 3 型号已可在 ROCm 下运行 PyTorch（示例：ROCm 5.7 起对 RX 7900 XT 提供 PyTorch 2.0.1 支持；Ubuntu 22.04.3 验证）。
  • 状态提醒：RDNA 4（如 RX 9070 系列）在发布时官方明确 ROCm 支持将“后续提供”，具体时间未定；ROCm 6.4 仍未包含 RDNA 4。
  • 建议：在目标 Ubuntu 版本上按 AMD 官方 ROCm 与 PyTorch 兼容矩阵安装，优先使用已验证的发行版与内核组合。

快速自检与选配步骤

• 查看 GPU 型号与驱动：nvidia-smi（NVIDIA）或 lspci | grep -i amd（AMD）。
• 查看计算能力（NVIDIA）：nvidia-smi -q | grep -i "Compute Capability"，或查阅 NVIDIA 官方 CC 表。
• 选择 PyTorch 构建：依据显卡 CC 与驱动支持的 CUDA 运行时，选择对应的 PyTorch + CUDA（或 ROCm）版本；优先用官方渠道预编译包。
• 验证：python - <<'PY' import torch print("CUDA可用性:", torch.cuda.is_available()) if torch.cuda.is_available(): print("设备:", torch.cuda.get_device_name(0)) print("计算能力:", torch.cuda.get_device_capability(0)) PY