性能优化指南

本文介绍如何优化深度学习训练性能，充分利用 GPU 算力。

GPU 利用率优化

监控 GPU 状态

# 实时监控
watch -n 1 nvidia-smi

# 查看详细信息
nvidia-smi -q

关注指标：

• GPU-Util：GPU 计算利用率，理想值 > 80%
• Memory-Usage：显存使用率
• Power：功耗，反映实际负载

提高 GPU 利用率

1. 增大 Batch Size

在显存允许范围内尽量增大 batch size：

# 找到最大可用 batch size
for batch_size in [16, 32, 64, 128, 256]:
    try:
        train_one_batch(batch_size)
        print(f"Batch size {batch_size} OK")
    except RuntimeError:
        print(f"Batch size {batch_size} OOM")
        break

2. 使用 DataLoader 预加载

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,      # 多进程加载
    pin_memory=True,    # 锁页内存，加速传输
    prefetch_factor=2,  # 预加载批次数
)

3. 避免 CPU-GPU 数据传输瓶颈

# 一次性将数据移到 GPU
data = data.to(device, non_blocking=True)

# 避免频繁的 .cpu() 和 .item() 调用
# 不好的做法
for batch in dataloader:
    loss = model(batch)
    print(loss.item())  # 每次都同步

# 好的做法
losses = []
for batch in dataloader:
    loss = model(batch)
    losses.append(loss)
# 最后统一处理

混合精度训练

使用 FP16 可以减少显存占用并加速计算：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()

    with autocast():
        output = model(data)
        loss = criterion(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

梯度累积

当显存不足以支持大 batch size 时：

accumulation_steps = 4
effective_batch_size = batch_size * accumulation_steps

optimizer.zero_grad()
for i, (data, target) in enumerate(dataloader):
    output = model(data)
    loss = criterion(output, target) / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

数据加载优化

使用 SSD 存储数据

将训练数据放在数据盘（本地 SSD）：

data_path = "/root/rivermind-data/datasets/imagenet"

预处理数据

提前处理好数据，避免训练时重复计算：

# 预处理并保存
processed_data = preprocess(raw_data)
torch.save(processed_data, "processed_data.pt")

# 训练时直接加载
data = torch.load("processed_data.pt")

使用内存映射

对于超大数据集：

import numpy as np

# 创建内存映射文件
data = np.memmap("data.bin", dtype="float32", mode="r", shape=(1000000, 256))

模型优化

使用 torch.compile（PyTorch 2.0+）

model = torch.compile(model)

启用 cuDNN 优化

torch.backends.cudnn.benchmark = True

禁用梯度计算（推理时）

with torch.no_grad():
    output = model(input)

多卡训练

DataParallel（简单）

model = torch.nn.DataParallel(model)

DistributedDataParallel（推荐）

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model, device_ids=[local_rank])

性能分析工具

PyTorch Profiler

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    model(input)

print(prof.key_averages().table(sort_by="cuda_time_total"))

简单计时

import time

torch.cuda.synchronize()
start = time.time()

# 训练代码
for epoch in range(10):
    train_one_epoch()

torch.cuda.synchronize()
print(f"耗时: {time.time() - start:.2f}s")

常见性能问题

问题	表现	解决方案
GPU 利用率低	GPU-Util < 50%	增大 batch size，优化数据加载
数据加载慢	GPU 等待数据	增加 num_workers，使用 SSD
显存不足	OOM 错误	减小 batch size，使用混合精度
训练速度慢	每步耗时长	使用 torch.compile，检查瓶颈