ubuntu pytorch怎么用gpu

在 Ubuntu 上使用 PyTorch 调用 GPU，一般按下面几步来：

一、确认环境

1. 确认有 NVIDIA 显卡

nvidia-smi

如果能看到显卡信息（如 CUDA Version），说明驱动已安装。

2. 查看系统 CUDA（驱动支持）

nvidia-smi

右上角会显示 CUDA Version（驱动支持的上限）
（不一定是系统已安装的 CUDA Toolkit）

二、安装 PyTorch（GPU 版）

推荐方式（CUDA 12.x，常见）

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

版本对应：

• cu118 → CUDA 11.8
• cu121 → CUDA 12.1
• cu124 → CUDA 12.4

可在官网查看： https://pytorch.org/get-started/locally/

三、验证 PyTorch 是否识别 GPU

import torch

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

四、在代码中使用 GPU

1. 自动选择设备

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

2. 模型放到 GPU

model = MyModel().to(device)

3. 数据放到 GPU

x = x.to(device)
y = y.to(device)

4. 训练示例

for x, y in dataloader:
    x, y = x.to(device), y.to(device)
    output = model(x)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()

五、常见问题

❌ torch.cuda.is_available() 返回 False

可能原因：

• 没装 GPU 版 PyTorch
• 驱动版本过低
• 系统 CUDA 与 PyTorch CUDA 不匹配（不影响，PyTorch 自带 CUDA）

✅ 检查 PyTorch 的 CUDA

print(torch.version.cuda)

六、Conda 用户（可选）

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

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

• Ubuntu 版本
• 显卡型号
• 是否已安装驱动

我可以给你精确命令。