如何设置PyTorch程序在 GPU上运行

1. 环境准备与检查 (系统兼容性)

Windows: 支持，但必须安装 NVIDIA 显卡驱动
Linux: 支持，是深度学习最推荐的系统
Mac (M系列芯片) : 支持，但使用的是 MPS (Metal Performance Shaders) 引擎，而非 CUDA,性能同样强劲
Mac (Intel芯片) : 不支持 GPU 加速
Windows (无独显/集显) : 不支持 CUDA，但可以安装 PyTorch 在 CPU 模式下运行

2. 检查驱动与 CUDA 版本 (关键步骤)

nvidia-smi命令:显示的是驱动支持的最高 CUDA 版本，而不是必须安装的版本

Windows/Linux: 打开终端 (CMD/PowerShell/Terminal)，输入：

nvidia-smi

图示说明 :
- 关注右上角的 CUDA Version: xx.x
- 只要这个版本 >= 11.8，就可以安装 PyTorch 的 CUDA 11.8 或 12.1 版本
- 如果提示 'nvidia-smi' 不是内部或外部命令，请安装 NVIDIA 显卡驱动

注意 : 不需要严格匹配 nvidia-smi 显示的 CUDA 版本,PyTorch 官网提供的 CUDA 11.8 或 CUDA 12.1 通常都能在较新的驱动上运行,如果驱动太老（比如 CUDA Version 显示 11.4 以下），则需要更新驱动

Mac (M系列: M1/M2/M3芯片**)**: 无需此步骤，直接安装 PyTorch 即可支持 MPS, 命令（无需 CUDA）如下:
python 复制代码
# 安装 PyTorch (支持 MPS)
pip3 install torch torchvision torchaudio

# 测试代码中使用 'mps' 而不是 'cuda'
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')

3. 获取安装命令 (官网交互式生成)

官网提供了配置生成器

1.访问 PyTorch 官网

2.点击顶部的 "Get Started"。

3.在 "Select Your Preferences" 区域，选择配置：

PyTorch Build: Stable (最新稳定版)

Your OS: 选择系统 (Windows, Linux, Mac)

Package: Pip (或 Conda，看习惯)

Language: Python

Compute Platform : 关键 ,如果有 NVIDIA 显卡，选择 CUDA (推荐 11.8 或 12.1)；如果是 Mac M芯片，选择 CPU (MPS 支持通常包含在 CPU 版本中)

4.复制下方生成的命令

4. 创建虚拟环境 (推荐)

使用 Conda 的方式，这里补充 Pip 的虚拟环境创建方式，更加通用

使用 Conda (推荐)

复制代码

1. conda create -n pytorch_gpu python=3.9
2. conda activate pytorch_gpu
3. 然后执行官网复制的安装命令

5. 测试代码 (验证安装)

在 PyCharm 或 Jupyter 中运行以下代码：

复制代码

import torch

# 检查 CUDA 是否可用 (NVIDIA 显卡)
if torch.cuda.is_available():
    print(" CUDA 可用！")
    print(f"GPU 数量: {torch.cuda.device_count()}")
    print(f"当前 GPU: {torch.cuda.get_device_name(0)}")
    # 将张量移动到 GPU
    x = torch.rand(3, 3).to('cuda') 
    print(x)
else:
    print(" CUDA 不可用")

# 针对 Mac M系列芯片的检查
if torch.backends.mps.is_available():
    print(" MPS 可用 (Mac M芯片)!")
    x = torch.rand(3, 3).to('mps')
    print(x)

python 复制代码

# 设备选择, 我们可以选择在cuda或者cpu上运行你的代码
# windows写法
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Mac写法
# device = "mps"
# print(f'当前设备: {device}')

6.常见问题与解决

"CUDA out of memory":

显存不足。尝试减小 Batch Size

或者在代码开头添加：torch.cuda.empty_cache()

"nvidia-smi works but PyTorch says CUDA not available":

这通常是因为安装了 CPU 版本的 PyTorch,请检查执行的安装命令是否包含 cu118 或 cu121,请务必去官网重新复制正确的命令

Mac M芯片速度慢:

确保代码中使用了 .to('mps')

目前 PyTorch 对 MPS 的支持仍在完善中，部分算子可能回退到 CPU，这是正常现象。

7.总结对比表

步骤	你的原方法	优化后的方法 (推荐)	优势
查看版本	`nvidia-smi`	`nvidia-smi` (仅看驱动是否支持)	不纠结具体版本号，只要驱动够新即可。
获取命令	手动去网页找 `cu118` 链接	官网交互式生成器	不容易出错，永远获取最新稳定命令
环境	Conda	Conda 或 Python venv	venv 更轻量，Conda 适合管理复杂依赖
Mac支持	未提及	明确区分 MPS 和 CPU	让 Mac 用户也能利用芯片算力