PyTorch GPU 配置指南
title: “PyTorch GPU 环境配置指南:CUDA 安装与验证” linkTitle: “GPU 配置” slug: “pytorch-gpu-setup” date: 2025-08-03T10:00:00+08:00 lastmod: 2026-08-18T14:07:00+08:00 draft: false description: “完整的 PyTorch GPU 环境配置教程：从 CUDA 版本验证到 GPU 版本安装，涵盖 NVIDIA 显卡驱动配置、AMD ROCm 环境搭建及常见问题排查。” categories:
- 人工智能
- 深度学习
- 环境配置 tags:
- PyTorch
- CUDA
- GPU 加速
- 深度学习框架
- 环境配置 author: “raikay” image: “https://blogimg.raikay.cn/images/20250803/PyTorch GPU 配置指南.png” weight: 2
PyTorch 作为主流的深度学习框架，在安装时严格区分了 CPU 版本与 GPU 版本。这两个版本对应不同的底层二进制分发包，并非仅通过代码参数进行区分。在实际开发与部署中，准确判断当前环境并进行相应的版本切换是保障模型训练与推理效率的基础。
一、验证当前 PyTorch 版本类型
确认已安装版本的最直接方式是通过 Python 代码进行检测。通过导入 torch 库并调用相关函数，可以获取当前的版本信息及硬件加速状态:
import torch
print("PyTorch 版本:", torch.__version__)
print("CUDA 是否可用:", torch.cuda.is_available())
print("CUDA 版本:", torch.version.cuda)
在输出结果中，若 torch.cuda.is_available() 返回 True,且版本号包含 +cuXXX 后缀 (如 2.1.0+cu118),则表明当前为 GPU 版本;若返回 False,或版本号包含 +cpu 后缀，则表明当前为纯 CPU 版本。此外，也可在终端使用 pip show torch 命令，通过查看包版本号的后缀来快速判断。
二、从 CPU 版本切换至 GPU 版本
若检测到当前为 CPU 版本且需要启用硬件加速，需执行卸载与重装操作。首先，通过 pip uninstall torch torchvision torchaudio -y 命令彻底卸载现有的 CPU 版本。
随后，需确认系统显卡驱动所支持的最高 CUDA 版本 (可通过终端执行 nvidia-smi 命令查看)。根据查到的 CUDA 版本，从 PyTorch 官方源拉取对应的 GPU 版本进行安装。例如，对于支持 CUDA 12.1 及以上版本的显卡，可执行以下命令:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
需要注意的是，GPU 版本的 PyTorch 必须从官方服务器下载，使用常规的国内镜像源通常只能获取到 CPU 版本。
三、特殊硬件架构的环境配置说明
PyTorch 的 GPU 加速生态不仅限于 NVIDIA 显卡。针对 AMD 显卡，官方提供了基于 ROCm (AMD 开源计算栈) 的适配版本。然而，ROCm 环境目前仅支持 Linux 操作系统 (如 Ubuntu)。在 Linux 环境下，用户需先安装 ROCm 系统依赖并配置环境变量，再通过 ROCm 专用的索引源安装 PyTorch。
对于在 Windows 系统下使用 AMD 显卡的用户，由于缺乏完整的 ROCm 支持，无法在原生 Windows 环境中启用 PyTorch 的 GPU 加速。此类场景下，通常需将代码中的设备参数强制指定为 CPU 运行，或通过配置 Linux 双系统、虚拟机等方式来搭建 ROCm 环境。