linux pytorch如何调GPU

在 Linux 下使用 PyTorch 调用 GPU，一般按下面几个步骤来：

一、确认环境

1. 确认有 NVIDIA GPU

nvidia-smi

如果能看到 GPU 信息，说明驱动已安装。

2. 确认 CUDA 可用

nvcc --version

或使用 PyTorch 检查：

import torch
print(torch.cuda.is_available())

返回 True 表示 PyTorch 能使用 GPU。

二、安装支持 GPU 的 PyTorch（Linux）

推荐使用官方命令（以 CUDA 11.8 为例）：

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

三、在代码中调用 GPU

1. 选择设备

import torch

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

2. 把模型放到 GPU

model = MyModel().to(device)

3. 把数据放到 GPU

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

四、多 GPU 情况（可选）

使用单卡指定

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

使用多卡（DataParallel，简单）

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

推荐方式（DistributedDataParallel）

适合训练大模型：

torchrun --nproc_per_node=2 train.py

五、常见坑

❌ torch.cuda.is_available() 为 False：

• 装了 CPU 版 PyTorch
• CUDA 驱动版本不匹配
• Linux 没装 NVIDIA 驱动

✅ 检查版本对应：

print(torch.__version__)
print(torch.version.cuda)

六、最简示例

import torch

x = torch.randn(3, 3).cuda()
y = torch.randn(3, 3).cuda()
z = x + y
print(z)

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

• 你的 GPU 型号
• CUDA 版本
• 是 训练 / 推理 / 服务器

我可以给你更具体的配置方案。