ubuntu怎么测试pytorch

在 Ubuntu 上测试 PyTorch 是否安装成功、能否正常使用（尤其是能否调用 GPU），可以按下面的步骤来。

一、确认 Python 和 pip

python3 --version
pip3 --version

建议使用虚拟环境（可选但推荐）：

python3 -m venv pytorch-env
source pytorch-env/bin/activate

二、安装 PyTorch（如尚未安装）

1️⃣ CPU 版本

pip install torch torchvision torchaudio

2️⃣ GPU 版本（CUDA，推荐）

去官网选命令：https://pytorch.org/get-started/locally/

例如（CUDA 11.8）：

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

三、测试 PyTorch 是否可用

1️⃣ 基本导入测试

python3 -c "import torch; print(torch.__version__)"

如果能输出版本号，说明安装成功。

2️⃣ 测试 CPU 是否可用

python3 -c "import torch; print(torch.cuda.is_available())"

• False → 只用 CPU
• True → GPU 可用

3️⃣ 完整测试脚本（推荐）

新建 test_pytorch.py：

import torch

print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("CUDA version:", torch.version.cuda)
    print("GPU name:", torch.cuda.get_device_name(0))

    x = torch.randn(3, 3).cuda()
    y = torch.randn(3, 3).cuda()
    print("GPU compute test:\n", x + y)
else:
    x = torch.randn(3, 3)
    y = torch.randn(3, 3)
    print("CPU compute test:\n", x + y)

运行：

python3 test_pytorch.py

四、常见问题排查

❌ torch.cuda.is_available() 为 False

可能原因：

• 没装 CUDA 版 PyTorch
• 驱动版本太低
• Ubuntu 没装 NVIDIA 驱动

检查驱动：

nvidia-smi

❌ 导入报错

pip uninstall torch
pip install torch

如果你愿意，可以告诉我：

• Ubuntu 版本
• 是否有 NVIDIA 显卡
• 是否用 conda / docker

我可以给你更精确的命令。