linux pytorch如何调用gpu

在 Linux 下使用 PyTorch 调用 GPU，通常需要完成环境准备、代码修改、验证三步。下面给你一个清晰、实用的说明。

一、确认 Linux 环境有 NVIDIA GPU

nvidia-smi

如果能看到 GPU 信息（如型号、显存、CUDA 版本），说明驱动已安装。

如果命令不存在，需要先安装 NVIDIA 驱动。

二、安装支持 GPU 的 PyTorch（CUDA 版）

1. 查看 CUDA 版本

nvidia-smi

右上角会显示 CUDA Version: 12.x

2. 安装 PyTorch（以 CUDA 12.1 为例）

推荐使用官方命令：

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

⚠️ 不要只用 pip install torch，默认可能是 CPU 版

三、在 PyTorch 中调用 GPU

1. 检查 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 = "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（可选）

model = torch.nn.DataParallel(model)

或使用更推荐的：

model = torch.compile(model)  # PyTorch 2.x

四、常见问题排查

❌ torch.cuda.is_available() 返回 False

可能原因：

• 安装了 CPU 版 PyTorch
• CUDA 驱动版本不匹配
• 系统没有 NVIDIA 驱动

解决：

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

❌ 显存不足（CUDA out of memory）

• 减小 batch_size
• 使用 torch.cuda.empty_cache()
• 使用混合精度：

with torch.cuda.amp.autocast():
    output = model(x)

五、最简示例

import torch

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

x = torch.randn(1000, 1000).to(device)
y = x @ x
print(y.device)

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

• 你的 Linux 发行版（Ubuntu / CentOS 等）
• GPU 型号
• CUDA 版本
• 是否已安装驱动

我可以给你一条完全对应你环境的安装命令。