在PyTorch中使用GPU加速：从基础操作到模型部署

本文将通过具体代码示例，详细介绍如何在PyTorch中利用GPU进行张量计算和模型训练，包含设备查询、数据迁移以及模型部署等完整流程。

---

1. 查看GPU硬件信息

使用 nvidia-smi 命令检查GPU状态和进程信息：

# 查看GPU信息
!nvidia-smi

输出示例：
显示当前GPU型号（如NVIDIA GeForce RTX 3060）、显存使用情况和运行中的进程。

---

2. 计算设备管理

2.1 定义设备对象

import torch
from torch import nn

# 定义不同设备对象
device_cpu = torch.device('cpu')
device_gpu0 = torch.cuda.device('cuda')
device_gpu1 = torch.cuda.device('cuda:1')
device_cpu, device_gpu0, device_gpu1

输出：

(device(type='cpu'), <torch.cuda.device at 0x...>, <torch.cuda.device at 0x...>)

2.2 查询可用GPU数量

torch.cuda.device_count()  # 返回可用GPU数量

输出：1（表示系统中有1块可用GPU）

3. 灵活选择设备

3.1 按需选择GPU或CPU

def try_gpu(i=0):
    """选择指定GPU，若不可用则返回CPU"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpu():
    """返回所有可用GPU，若无则返回CPU"""
    devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

try_gpu(), try_gpu(10), try_all_gpu()

 输出：

(device(type='cuda', index=0), device(type='cpu'), [device(type='cuda', index=0)])

4. 张量与GPU操作

4.1 默认设备查询

x = torch.tensor([1, 2, 3])
x.device  # 默认在CPU上

输出：device(type='cpu')

4.2 显存中创建张量

x_gpu = torch.ones(2, 3, device='cuda:0')  # 在GPU 0上创建全1张量
x_gpu

输出：

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

4.3 GPU间计算

y_gpu = torch.rand(2, 3, device='cuda:0')
x_gpu + y_gpu  # 必须同一设备才能计算

输出示例：

tensor([[1.4571, 1.4172, 1.1364],
        [1.7572, 1.3977, 1.2402]], device='cuda:0')

5. 神经网络模型部署到GPU

5.1 模型迁移

# 定义一个简单神经网络
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device='cuda')  # 将模型参数移至GPU

# 输入数据需与模型在同一设备
output = net(x_gpu)
output

输出示例：

tensor([[-0.4271],
        [-0.4271]], device='cuda:0', grad_fn=<AddmmBackward0>)

5.2 验证参数位置

net[0].weight.data.device  # 检查参数存储位置

输出：device(type='cuda', index=0)

6. 关键注意事项

1. 设备一致性：输入数据和模型必须在同一设备（CPU/GPU）上才能进行计算。

2. 显存管理：及时释放不再使用的GPU张量（del tensor + torch.cuda.empty_cache()）。

3. 多GPU支持：可通过torch.nn.DataParallel实现多卡并行训练。

---

总结

本文演示了PyTorch中GPU加速的核心操作，包括设备选择、张量迁移和模型部署。合理利用GPU可显著加速深度学习任务，建议在训练大型模型时优先使用GPU环境。如果遇到CUDA相关错误，请检查驱动版本和PyTorch的GPU支持安装。