🚀 GPU加速训练

GPU可以大幅加速深度学习训练。本节介绍如何在PyTorch中使用GPU。

📌 支持的加速设备

PyTorch 支持多种硬件加速：

NVIDIA CUDA - 最常用，Windows/Linux均支持
AMD ROCm - 仅Linux，API与CUDA兼容
Apple MPS - Apple Silicon (M1/M2/M3) 专用

🤔 为什么用GPU？

操作	CPU	GPU
矩阵乘法	顺序执行	并行执行（几千个核心）
训练速度	慢	快10-100倍
适合场景	小模型、推理	大模型、训练

💡 简单理解

CPU：像一个高学历博士，什么都会，但只有几个人
GPU：像几千个高中生，只会简单计算，但人多力量大
深度学习需要大量简单的矩阵运算，所以GPU更合适

✅ 检查GPU是否可用

import torch

# ============ NVIDIA CUDA ============
print(f"CUDA可用: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"CUDA版本: {torch.version.cuda}")
    print(f"GPU数量: {torch.cuda.device_count()}")
    print(f"GPU名称: {torch.cuda.get_device_name(0)}")
    print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

# ============ Apple MPS (M1/M2/M3) ============
print(f"MPS可用: {torch.backends.mps.is_available()}")

# ============ AMD ROCm ============
# ROCm使用相同的API: torch.cuda.is_available()

📱 设备管理

定义设备（推荐方式）

根据 PyTorch官方文档推荐的设备选择模式：

import torch

# 推荐：自动选择最佳设备
device = (
    "cuda" if torch.cuda.is_available()       # NVIDIA GPU 或 AMD ROCm
    else "mps" if torch.backends.mps.is_available()  # Apple Silicon
    else "cpu"
)
print(f"使用设备: {device}")

其他设备定义方式

import torch

# 方式1：指定GPU编号（多卡情况）
device = torch.device('cuda:0')  # 第一张GPU
device = torch.device('cuda:1')  # 第二张GPU

# 方式2：直接使用字符串
device = 'cuda'  # PyTorch会自动转换

# 方式3：Apple MPS
device = torch.device('mps')  # Apple Silicon

移动数据到GPU

import torch

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

# 创建张量时直接指定设备
x = torch.randn(3, 4, device=device)
print(x.device)  # cuda:0

# 将已有张量移到GPU
y = torch.randn(3, 4)
y = y.to(device)
# 或
y = y.cuda()  # 移到默认GPU

# 移回CPU
z = y.cpu()
# 或
z = y.to('cpu')

移动模型到GPU

import torch.nn as nn

model = nn.Linear(10, 5)

# 移动到GPU
model = model.to(device)
# 或
model = model.cuda()

# 检查模型在哪个设备
print(next(model.parameters()).device)  # cuda:0

🔧 完整的GPU训练流程

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# ============ 1. 设置设备 ============
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# ============ 2. 准备数据 ============
X = torch.randn(1000, 20)
y = torch.randint(0, 2, (1000,))
dataset = TensorDataset(X, y)

# pin_memory=True 可以加速CPU到GPU的数据传输
train_loader = DataLoader(
    dataset, 
    batch_size=32, 
    shuffle=True,
    pin_memory=True  # GPU训练时启用
)

# ============ 3. 定义模型并移到GPU ============
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(20, 64)
        self.fc2 = nn.Linear(64, 2)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = Net().to(device)  # 移到GPU

# ============ 4. 损失函数和优化器 ============
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# ============ 5. 训练循环 ============
for epoch in range(10):
    model.train()
    for data, labels in train_loader:
        # 数据移到GPU
        data = data.to(device)
        labels = labels.to(device)
        
        # 训练步骤
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch+1} 完成")

print("训练完成！")

⚠️ 常见错误

错误1：数据和模型不在同一设备

# ❌ 错误
model = model.cuda()
x = torch.randn(10, 20)  # 在CPU上
output = model(x)  # 报错！

# ✅ 正确
model = model.cuda()
x = torch.randn(10, 20).cuda()
output = model(x)  # 正常

错误2：GPU上的张量无法直接转NumPy

# ❌ 错误
x = torch.randn(3, 4).cuda()
np_array = x.numpy()  # 报错！

# ✅ 正确
x = torch.randn(3, 4).cuda()
np_array = x.cpu().numpy()  # 先移到CPU

错误3：忘记移动数据

# ❌ 错误：忘记把labels也移到GPU
data = data.to(device)
# labels还在CPU上
loss = criterion(outputs, labels)  # 可能报错

# ✅ 正确
data = data.to(device)
labels = labels.to(device)

📊 监控GPU使用

在代码中查看

import torch

# 查看当前GPU内存使用
print(f"已分配内存: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"缓存内存: {torch.cuda.memory_reserved() / 1e9:.2f} GB")

# 清空缓存
torch.cuda.empty_cache()

使用命令行

# 实时监控
nvidia-smi

# 每秒刷新
watch -n 1 nvidia-smi

🎯 GPU内存优化技巧

1. 减小batch_size

# 内存不足时，减小batch_size
train_loader = DataLoader(dataset, batch_size=16)  # 从32减到16

2. 使用混合精度训练

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, labels in train_loader:
    data, labels = data.to(device), labels.to(device)
    
    optimizer.zero_grad()
    
    # 使用混合精度
    with autocast():
        outputs = model(data)
        loss = criterion(outputs, labels)
    
    # 缩放梯度
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

3. 梯度累积

# 模拟大batch_size
accumulation_steps = 4
optimizer.zero_grad()

for i, (data, labels) in enumerate(train_loader):
    data, labels = data.to(device), labels.to(device)
    
    outputs = model(data)
    loss = criterion(outputs, labels) / accumulation_steps
    loss.backward()
    
    # 每accumulation_steps步更新一次
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

4. 推理时释放梯度

# 推理时不需要梯度
model.eval()
with torch.no_grad():
    outputs = model(data)

🖥️ 多GPU训练（简介）

如果你有多张GPU，可以使用DataParallel：

import torch.nn as nn

# 单行代码实现多GPU
if torch.cuda.device_count() > 1:
    print(f"使用 {torch.cuda.device_count()} 张GPU")
    model = nn.DataParallel(model)

model = model.to(device)

💡 提示

更高级的多GPU训练可以使用DistributedDataParallel，效率更高，但配置更复杂。

📝 最佳实践

始终检查设备：用torch.cuda.is_available()
统一设备管理：定义一个device变量，所有地方用它
启用pin_memory：DataLoader中设置pin_memory=True（仅CUDA）
监控内存：避免OOM（内存溢出）错误
合理batch_size：GPU内存允许的情况下尽量大

🍎 Apple Silicon (MPS) 用户指南

如果你使用的是 Mac M1/M2/M3 芯片，可以使用 MPS 加速：

import torch

# 检查MPS是否可用
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print("使用 Apple MPS 加速")
else:
    device = torch.device("cpu")
    print("MPS不可用，使用CPU")

# 创建张量并移到MPS
x = torch.randn(3, 3, device=device)
print(x)

# 模型移到MPS
model = model.to(device)

⚠️ MPS注意事项

MPS 不支持所有 CUDA 操作，部分高级功能可能不可用
pin_memory=True 在 MPS 上无效，可能导致警告
建议使用 PyTorch 2.0+ 以获得更好的 MPS 支持

📚 官方资源

🎉 进阶篇完成！

恭喜你完成了进阶篇的学习！现在你已经能够：

构建神经网络
训练模型
保存和加载模型
使用GPU加速（CUDA/MPS）

接下来，让我们进入高级篇，学习CNN、RNN等高级网络！