在 PyTorch 中使用 GPU 加速模型训练是深度学习的常见需求，以下是关键方法和步骤的中文总结：

1. 确保 GPU 可用性

import torch

# 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前设备: {device}")

2. 将模型和数据移动到 GPU

模型迁移到 GPU

model = YourModel().to(device)  # 将模型参数和缓冲区移至 GPU

数据迁移到 GPU

# 单样本数据
inputs = inputs.to(device)
labels = labels.to(device)

# 批量数据（推荐在 DataLoader 中直接处理）

3. 使用 DataLoader 加速数据加载

通过 DataLoader 配置高效的数据加载：

from torch.utils.data import DataLoader

dataloader = DataLoader(dataset, 
                        batch_size=64, 
                        shuffle=True,
                        num_workers=4,     # 多进程加载数据（根据 CPU 核心数调整）
                        pin_memory=True)   # 启用锁页内存，加速 GPU 传输

4. 混合精度训练（AMP）

使用 torch.cuda.amp 自动混合精度，减少显存占用并加速计算：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, labels in dataloader:
    inputs, labels = inputs.to(device), labels.to(device)

    optimizer.zero_grad()

    with autocast():           # 自动转换精度
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)

    scaler.scale(loss).backward()  # 缩放梯度
    scaler.step(optimizer)         # 更新参数
    scaler.update()                # 调整缩放因子

5. 多 GPU 并行训练

DataParallel（单机多卡，简单但效率较低）

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

DistributedDataParallel（DDP，推荐用于多机多卡）

# 初始化进程组
torch.distributed.init_process_group(backend="nccl")
model = nn.parallel.DistributedDataParallel(model)

6. 优化 GPU 计算效率

启用 cuDNN 加速

torch.backends.cudnn.benchmark = True   # 自动寻找最优卷积算法（输入尺寸固定时启用）

避免 CPU-GPU 频繁传输

• 数据预处理尽量放在 CPU 或 DataLoader 中。
• 减少训练循环中的 print 或日志操作。

7. 梯度累积（显存不足时使用）

通过累积多个小批次的梯度模拟大 batch size：

accumulation_steps = 4  # 累积 4 个 batch 的梯度

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = loss_fn(outputs, labels)
    loss = loss / accumulation_steps  # 梯度归一化
    loss.backward()

    if (i+1) % accumulation_steps == 0:
        optimizer.step()     # 更新参数
        optimizer.zero_grad()  # 清空梯度

8. 监控 GPU 使用情况

• 终端命令：nvidia-smi 查看显存占用和 GPU 利用率。
• PyTorch 内置工具：

  print(torch.cuda.memory_allocated())  # 当前显存占用量
  print(torch.cuda.max_memory_allocated())  # 历史最大占用量

注意事项

1. 版本匹配：确保 PyTorch、CUDA 和 cuDNN 版本兼容。
2. 数据平衡：避免单个 GPU 负载不均（多卡训练时）。
3. 释放显存：torch.cuda.empty_cache() 可手动清理缓存碎片。

合理组合这些方法，可显著提升 PyTorch 模型的训练速度！