来自 千问
要让 PyTorch 模型在 GPU 上高效运行，并不仅仅是改变设备设置，更是一套涉及环境、代码和最佳实践的组合拳。简单来说，这包括三个核心方面：确保环境支持、正确编写代码、以及应用性能优化技巧。
以下是对如何在 GPU 上进行模型训练和测试的全面总结与示例。
🚀 核心要点：如何确保"尽量在 GPU 上跑"？
-
环境是基础
- 驱动与 CUDA：首先，你的电脑需要安装 NVIDIA 显卡驱动和 CUDA Toolkit。这是 GPU 能够被调用的基础环境。
- PyTorch 版本 ：必须安装支持 CUDA 的 PyTorch 版本。可以通过
torch.cuda.is_available()来快速验证环境是否配置正确。
-
代码是关键
- 统一设备：模型、输入数据、标签这三者必须都在同一个设备（GPU）上。如果模型在 GPU，而标签在 CPU，程序会立即报错。
- 不要遗漏：除了模型和输入数据，计算损失函数（Loss Function）时用到的**标签（Labels）**也必须移动到 GPU。
-
优化是进阶
- 混合精度 ：使用
torch.cuda.amp可以让计算在float16半精度下进行，这能显著减少显存占用，并提升计算速度。 - 数据加载 ：在
DataLoader中设置pin_memory=True，可以加速数据从 CPU 内存传输到 GPU 显存的过程。
- 混合精度 ：使用
💻 完整示例代码
这是一个基于 CIFAR-10 数据集的完整训练与测试流程示例，包含了上述所有最佳实践。
python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# 1. 设置设备
# 自动选择可用的 GPU，如果没有则回退到 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"🚀 正在使用设备: {device}")
# 2. 数据准备
# 定义图像预处理流程
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 标准化到 [-1, 1]
])
# 加载 CIFAR-10 数据集
train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
# 创建数据加载器
# pin_memory=True 可以加速 GPU 的数据传输
train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_set, batch_size=128, shuffle=False, num_workers=2, pin_memory=True)
# 3. 定义模型
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(64 * 8 * 8, 512),
nn.ReLU(),
nn.Linear(512, 10)
)
def forward(self, x):
x = self.features(x)
x = self.classifier(x)
return x
# 实例化模型并移动到 GPU
model = SimpleCNN().to(device)
# 4. 配置优化器与损失函数
criterion = nn.CrossEntropyLoss() # 损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001) # 优化器
# 5. (关键) 混合精度训练配置
# 这是提升 GPU 训练速度和减少显存占用的关键技术
scaler = torch.cuda.amp.GradScaler(enabled=(device.type == 'cuda'))
# 6. 训练与测试函数
def train(epoch):
model.train() # 切换到训练模式
for batch_idx, (data, target) in enumerate(train_loader):
# 🔥 核心操作：将数据和标签移动到 GPU
# non_blocking=True 允许异步数据传输，进一步提升效率
data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
optimizer.zero_grad()
# 使用混合精度上下文管理器
with torch.cuda.amp.autocast(enabled=(device.type == 'cuda')):
output = model(data)
loss = criterion(output, target)
# 反向传播 (使用缩放器处理梯度)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
if batch_idx % 100 == 0:
print(f"Train Epoch: {epoch} [{batch_idx}/{len(train_loader)}] Loss: {loss.item():.6f}")
def test():
model.eval() # 切换到评估模式
correct = 0
total = 0
with torch.no_grad(): # 测试时不需要计算梯度
for data, target in test_loader:
# 🔥 核心操作：将测试数据和标签也移动到 GPU
data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
output = model(data)
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
print(f"Test Accuracy: {100.*correct/total:.2f}%")
# 7. 开始训练
if __name__ == '__main__':
for epoch in range(1, 6): # 训练 5 个 epoch
train(epoch)
test()
# 8. 保存模型 (推荐只保存模型参数)
torch.save(model.state_dict(), "cifar10_model.pth")
print("模型训练完成并已保存！")📌 总结
通过上述代码，你实现了一个完整的、针对 GPU 优化的训练流程。关键点在于：
- 环境验证 ：确保
torch.cuda.is_available()返回True。 - 设备一致性：模型、数据、标签三者缺一不可，必须同在 GPU 上。
- 性能优化 ：利用
pin_memory、non_blocking和amp(混合精度) 等技术，让 GPU 的计算能力得到最大程度的发挥，避免数据传输成为瓶颈。