本文将通过具体代码示例，详细介绍如何在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. 关键注意事项
-
设备一致性：输入数据和模型必须在同一设备（CPU/GPU）上才能进行计算。
-
显存管理：及时释放不再使用的GPU张量（
del tensor+torch.cuda.empty_cache()）。 -
多GPU支持：可通过
torch.nn.DataParallel实现多卡并行训练。
总结
本文演示了PyTorch中GPU加速的核心操作，包括设备选择、张量迁移和模型部署。合理利用GPU可显著加速深度学习任务，建议在训练大型模型时优先使用GPU环境。如果遇到CUDA相关错误，请检查驱动版本和PyTorch的GPU支持安装。
)
)
)