PyTorch 性能分析与调优

性能分析是优化的第一步。PyTorch Profiler 提供了详细的执行时间、内存和 Kernel 分析能力。

Profiler 使用

from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
    record_shapes=True,
    profile_memory=True
) as prof:
    for i, batch in enumerate(dataloader):
        with record_function("model_forward"):
            output = model(batch)
        with record_function("model_backward"):
            loss.backward()
        prof.step()

# 查看热点
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

常见性能瓶颈

1. 数据加载慢 → 增加 num_workers
2. GPU 利用率低 → 增大 batch_size
3. 内存碎片 → 使用 torch.cuda.empty_cache()
4. Kernel 效率低 → 使用 torch.compile()

相关资源

• PyTorch Profiler 文档
• TensorBoard
• PyTorch 性能调优指南