ubuntu pytorch如何调用GPU

在 Ubuntu 上使用 PyTorch 调用 GPU，通常步骤如下：

一、确认环境

1. 确认有 NVIDIA GPU

nvidia-smi

能看到显卡信息说明驱动已安装。

2. 查看 CUDA 版本

nvidia-smi

右上角会显示 CUDA Version（驱动支持的最大版本）

二、安装支持 GPU 的 PyTorch

推荐方式（官方命令）

去 https://pytorch.org 选择对应配置，例如：

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

cu118 表示 CUDA 11.8，按你的驱动版本选择

三、在代码中调用 GPU

1. 检查 PyTorch 是否能识别 GPU

import torch

print(torch.cuda.is_available())   # True 表示可用
print(torch.cuda.device_count())   # GPU 数量
print(torch.cuda.get_device_name(0))

2. 将模型和数据放到 GPU

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = MyModel().to(device)

x = torch.randn(8, 3, 224, 224).to(device)
y = model(x)

3. 指定某块 GPU（多卡情况）

device = torch.device("cuda:0")

四、常见问题

❌ torch.cuda.is_available() 为 False

可能原因：

• 装了 CPU 版 PyTorch
• 驱动未安装或不兼容
• CUDA 版本不匹配

解决：

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

✅ 多 GPU 训练（简单方式）

model = torch.nn.DataParallel(model)
model.to(device)

（更推荐 torch.distributed）

五、示例完整代码

import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = nn.Linear(10, 2).to(device)
x = torch.randn(32, 10).to(device)

out = model(x)
print(out.device)

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

• Ubuntu 版本
• GPU 型号
• 是否已安装 CUDA

我可以给你精确安装命令 ✅