Python代码怎么用GPU跑：CUDA环境配置与可用性验证

巷口听风

发布时间：2026-10-25| 4人浏览过

原创

让Python代码在GPU上运行，核心步骤是三件事：安装与显卡匹配的NVIDIA驱动，配置对应版本的CUDA Toolkit，再安装支持GPU的Python深度学习框架（如PyTorch或TensorFlow）。完成这三步后，用一段简短的Python代码即可验证GPU是否真正可用。本文按照这个顺序，把每一步需要做什么、怎么确认做对了讲清楚。

一、GPU跑Python代码的前提条件

并非所有显卡都能用来加速Python计算。要使用CUDA，需要满足以下硬件和软件条件：

NVIDIA显卡：CUDA是NVIDIA的专有技术，AMD显卡和Intel集成显卡不适用。笔记本和台式机的GeForce系列、工作站的Quadro系列、服务器的Tesla/A100/H100系列均可。
计算能力（Compute Capability）：显卡需要具备一定的计算能力版本。当前主流框架通常要求计算能力 ≥ 3.5，近两三年的显卡（如RTX 30系列、40系列）均满足。可在NVIDIA官网查询具体型号对应的计算能力。
操作系统：Windows 10/11、主流Linux发行版（Ubuntu、CentOS等）均支持。macOS从CUDA 11起已不再支持。

确认显卡型号最直接的方式：Windows在"设备管理器 → 显示适配器"中查看；Linux执行 lspci | grep -i nvidia。

二、安装NVIDIA驱动

驱动是操作系统与GPU通信的基础，必须先装好驱动，CUDA才能工作。

Windows

访问 https://www.nvidia.com/drivers，选择显卡型号和操作系统版本，下载最新的Game Ready或Studio驱动，运行安装程序，按默认选项完成即可。安装完成后打开命令提示符，执行：

nvidia-smi

如果输出一个表格，显示了驱动版本（Driver Version）和支持的CUDA版本（CUDA Version），说明驱动安装成功。示例输出：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.129.03   Driver Version: 535.129.03   CUDA Version: 12.2    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M.  |
|===============================+======================+======================|
|   0  NVIDIA GeForce RTX 4090   Off  | 00000000:01:00.0  On |                  N/A |
|  0%   35C    P8    20W / 450W |    512MiB / 24564MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

这里的"CUDA Version: 12.2"表示该驱动最高支持CUDA 12.2，后面安装CUDA Toolkit时版本不能超过这个数字。

Linux（Ubuntu）

Ubuntu推荐使用系统包管理器安装驱动，避免手动下载 .run 文件引起的兼容问题：

sudo apt update
sudo apt install nvidia-driver-535

版本号根据实际需求选择。安装完成后重启系统，再执行 nvidia-smi 确认。

注意：nvidia-smi 输出中的 CUDA Version 是驱动支持的最高CUDA版本，不代表系统已经安装了CUDA Toolkit。驱动和CUDA Toolkit是两个独立的组件。

图1：NVIDIA驱动安装与验证流程——先确认显卡型号，安装驱动后通过nvidia-smi验证

三、安装CUDA Toolkit

CUDA Toolkit包含编译器（nvcc）、运行时库和开发头文件。深度学习框架的GPU版本在底层依赖这些组件。

选择版本

CUDA Toolkit版本必须不超过驱动支持的CUDA版本。例如 nvidia-smi 显示"CUDA Version: 12.2"，则可安装CUDA 12.2或更低版本。同时，还需要查看目标框架所要求的CUDA版本：

框架	推荐CUDA版本	查看方式
PyTorch 2.x	CUDA 11.8 或 12.1/12.4	pytorch.org 安装页面
TensorFlow 2.14+	CUDA 12.3	tensorflow.org 安装指南
TensorFlow 2.10-2.13	CUDA 11.8	tensorflow.org 安装指南

框架的实际版本要求会随更新变化，安装前务必查阅对应官网的安装说明页面。

下载与安装

前往 https://developer.nvidia.com/cuda-toolkit-archive，选择目标版本和操作系统，按页面指引下载。

Windows：下载 .exe 安装包，运行后选择"自定义安装"，勾选CUDA组件（Runtime、Development、Documentation），按默认路径安装。安装完成后，系统环境变量中会自动添加 CUDA_PATH，同时 nvcc 会被加入 PATH。

Linux（Ubuntu，以CUDA 12.2为例）：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
sudo apt install cuda-toolkit-12-2

安装完成后，需要手动将CUDA加入环境变量。编辑 ~/.bashrc（或 ~/.zshrc），添加：

export PATH=/usr/local/cuda-12.2/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.2/lib64:$LD_LIBRARY_PATH

保存后执行 source ~/.bashrc 使其生效。

验证CUDA Toolkit

在终端执行：

nvcc --version

正常输出示例：

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Aug_15_22:02:13_PDT_2023
Cuda compilation tools, release 12.2, V12.2.140
Build cuda_12.2.r12.2/compiler.33191640_0

如果提示"command not found"，说明PATH未正确配置或CUDA未安装成功，需要回头检查。

四、安装cuDNN（可选但推荐）

cuDNN是NVIDIA针对深度神经网络优化的加速库，TensorFlow和PyTorch的很多GPU运算会调用它。没有cuDNN，部分操作可能无法使用GPU加速，或运行速度明显偏慢。

前往 https://developer.nvidia.com/cudnn（需注册NVIDIA开发者账号），下载与CUDA版本匹配的cuDNN。Linux下用tar包安装的操作步骤：

tar -xvf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz
sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/include/* /usr/local/cuda-12.2/include/
sudo cp cudnn-linux-x86_64-8.9.7.29_cuda12-archive/lib/* /usr/local/cuda-12.2/lib64/
sudo ldconfig

Windows下，将cuDNN解压后的 bin、include、lib 文件夹中的文件分别复制到CUDA安装目录的同名文件夹中即可。

图2：CUDA环境各组件的层次关系——从GPU硬件到Python代码的完整调用链

五、安装支持GPU的Python框架

驱动和CUDA Toolkit就绪后，接下来安装Python框架的GPU版本。强烈建议在虚拟环境中安装，避免与系统Python或其他项目冲突。

创建虚拟环境

python -m venv gpu_env
# Windows激活：
gpu_env\Scripts\activate
# Linux/macOS激活：
source gpu_env/bin/activate

安装PyTorch（GPU版本）

PyTorch的GPU版本自带对应的CUDA运行时，即使系统没有单独安装CUDA Toolkit，也能工作（前提是驱动已安装）。在 pytorch.org 首页选择操作系统、包管理器和CUDA版本后，会生成对应的安装命令。例如CUDA 12.4：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

安装完成后，文件体积较大（约2-3GB），耐心等待。

安装TensorFlow（GPU版本）

TensorFlow 2.x在2.10版本之后，pip安装的包已默认包含GPU支持，不需要单独安装 tensorflow-gpu。但TensorFlow要求系统已安装匹配版本的CUDA Toolkit和cuDNN。

pip install tensorflow

关键区别：PyTorch的pip包自带CUDA运行时库，装好驱动就能用；TensorFlow需要系统级安装CUDA Toolkit和cuDNN。这是两个框架在GPU环境配置上最大的不同。

六、验证GPU可用性

所有组件安装完成后，用Python代码验证GPU是否真正可用。

PyTorch验证

import torch

print("PyTorch版本:", torch.__version__)
print("CUDA是否可用:", torch.cuda.is_available())
print("CUDA版本:", torch.version.cuda)
print("GPU数量:", torch.cuda.device_count())
if torch.cuda.is_available():
    print("GPU名称:", torch.cuda.get_device_name(0))
    # 实际运算测试
    x = torch.randn(3, 3).cuda()
    print("GPU上的张量:\n", x)
    print("张量所在设备:", x.device)

期望输出（示例）：

PyTorch版本: 2.1.0+cu121
CUDA是否可用: True
CUDA版本: 12.1
GPU数量: 1
GPU名称: NVIDIA GeForce RTX 4090
GPU上的张量:
 tensor([[ 0.5432, -1.2301,  0.8723],
        [-0.1234,  0.9876, -0.5678],
        [ 0.3456, -0.7890,  0.1234]], device='cuda:0')
张量所在设备: cuda:0

关键看两点：torch.cuda.is_available() 返回 True，张量的device显示 cuda:0。

TensorFlow验证

import tensorflow as tf

print("TensorFlow版本:", tf.__version__)
print("GPU设备列表:", tf.config.list_physical_devices('GPU'))

# 实际运算测试
if tf.config.list_physical_devices('GPU'):
    with tf.device('/GPU:0'):
        a = tf.constant([[1.0, 2.0], [3.0, 4.0]])
        b = tf.constant([[5.0, 6.0], [7.0, 8.0]])
        c = tf.matmul(a, b)
        print("GPU矩阵乘法结果:\n", c.numpy())

期望输出（示例）：

TensorFlow版本: 2.15.0
GPU设备列表: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
GPU矩阵乘法结果:
 [[19. 22.]
 [43. 50.]]

如果GPU设备列表为空，说明TensorFlow没有检测到可用的GPU，需要检查CUDA和cuDNN的版本匹配。

图3：GPU可用性验证的逐步排查流程——依次检查驱动、CUDA Toolkit、框架GPU支持

七、常见问题与解决方法

torch.cuda.is_available() 返回 False

这是最常见的问题，依次检查以下几点：

驱动未安装或版本过低：执行 nvidia-smi，确认驱动正常且支持的CUDA版本 ≥ 框架要求的CUDA版本。
安装了CPU版本的PyTorch：如果安装时没有指定CUDA版本的 --index-url，pip默认安装CPU版本。用 pip show torch 查看版本号，如果版本号中不包含 cu（如 2.1.0+cpu），需要卸载后重新安装GPU版本。
虚拟环境混乱：确认当前激活的虚拟环境是安装了GPU版本框架的环境。用 which python（Linux）或 where python（Windows）确认Python路径。

TensorFlow检测不到GPU

TensorFlow对CUDA和cuDNN的版本要求更严格。TF 2.15需要CUDA 12.3和cuDNN 8.9，版本号不匹配时GPU不会被识别。可以通过以下命令查看TensorFlow的GPU日志：

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0'
import tensorflow as tf

日志中如果出现类似"Could not load dynamic library 'libcudart.so.12'"的信息，说明缺少对应版本的CUDA运行时库。

CUDA out of memory

GPU显存不足时会报此错误。减小batch size是最直接的解决办法。也可以用 torch.cuda.empty_cache() 释放未使用的缓存（注意这不会释放正在使用的显存）。用 nvidia-smi 查看是否有其他进程占用了显存。

八、将代码迁移到GPU上运行

验证环境没问题后，在实际代码中使用GPU需要注意的关键操作：

PyTorch中使用GPU

import torch

# 选择设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 将模型和数据转移到GPU
model = MyModel().to(device)
data = data.to(device)
labels = labels.to(device)

# 训练循环中，所有参与运算的张量都需要在同一设备上
output = model(data)
loss = criterion(output, labels)

关键原则：模型和所有参与运算的张量必须在同一个设备上。如果模型在GPU上，输入数据也必须先 .to(device) 转移到GPU。

TensorFlow中使用GPU

TensorFlow会自动将运算分配到可用的GPU上，通常不需要手动指定设备。如果需要控制：

import tensorflow as tf

# 查看可用GPU
gpus = tf.config.list_physical_devices('GPU')

# 限制GPU显存按需增长，避免一次性占满
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# 指定使用某个GPU
with tf.device('/GPU:0'):
    model = tf.keras.Sequential([...])
    model.compile(...)
    model.fit(x_train, y_train, ...)

图4：完整配置流程总结——PyTorch和TensorFlow两条路线的步骤差异

九、环境配置速查

将上述所有验证命令汇总，方便快速排查环境状态：

# 1. 检查驱动和GPU状态
nvidia-smi

# 2. 检查CUDA Toolkit版本
nvcc --version

# 3. 检查Python环境中的框架
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
# 或
python -c "import tensorflow as tf; print(tf.__version__, tf.config.list_physical_devices('GPU'))"

# 4. 检查cuDNN版本（PyTorch环境）
python -c "import torch; print(torch.backends.cudnn.version())"

如果第1步失败，问题在驱动；第1步通过但第2步失败，检查CUDA Toolkit安装和PATH配置；前两步都通过但第3步返回False或空列表，说明框架版本与CUDA版本不匹配，需要重新安装匹配的框架版本。按这个顺序逐步排查，可以定位到绝大多数GPU环境问题。