1. 检查是否安装了 GPU 版本的框架

TensorFlow

import tensorflow_x as tf

print(tf.config.list_physical_devices('GPU'))  # 输出 GPU 信息
  • 若返回空列表:
    • 安装 GPU 版本:pip install tensorflow-gpu(旧版本)或 pip install tensorflow[and-cuda](新版本)。
    • 确认 CUDA/cuDNN 版本与 TensorFlow 版本匹配(官方文档)。

PyTorch

import torch
print(torch.cuda.is_available())  # 输出 True 表示 GPU 可用

若返回 False:
安装支持 GPU 的 PyTorch(从官网选择对应 CUDA 版本的安装命令)。
示例:pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

查看CUDA版本

nvidia-smi

Logo

分享最新的 NVIDIA AI Software 资源以及活动/会议信息,精选收录AI相关技术内容,欢迎大家加入社区并参与讨论。

更多推荐