【PyTorch】使用Torch自带的工具进行性能分析

I. 常见工具介绍

1.1 Torch API

Torch.utils.bottleneck：脚本瓶颈分析工具

正如Torch官方文档说到：“torch.utils.bottleneck is a tool that can be used as an initial step for debugging bottlenecks in your program. It summarizes runs of your script with the Python profiler and PyTorch’s autograd profiler.”

Torch.profiler：性能分析新版API

Torch.profiler是替代torch.autograd.profiler的新版性能分析工具，可以结合FlameGraph和TensorBoard来可视化分析结果；

Torch.autograd.profiler (legacy)

关于之前使用torch.autograd.profiler常见的测试方法，请参考《pytorch中分析时间开销的方法》

1.2 Profile Visualization

Flame Graphs

Speedscope: an interactive FlameGraph visualizer

II. 性能分析步骤

PyTorch官方教程：《PYTORCH PROFILER》

7. 查看性能分析记录

导出性能分析的堆栈信息（.txt文件），

prof.export_stacks("/tmp/profiler_stacks.txt", "self_cuda_time_total")

8. 分析运行时间长的任务

8.1 常用组件

• schedule：用于自定义记录时间表，使用此参数的最好方法是基于torch.profiler.schedule()辅助函数，可以设置具体的日志时间表。关于 skip_first, wait, warmup, active, repeat参数的执行过程，请参考《PyTorch Profiler | example with torch.profiler.schedule helper function》

def trace_handler(p):
    output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)
    print(output)
    p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    # 设置记录时间表
    schedule=torch.profiler.schedule(
        wait=1,
        warmup=1,
        active=2),
    on_trace_ready=trace_handler
) as p:
    for idx in range(iters):
        model(inputs)
        p.step()

8.2 在TensorBoard中可视化堆栈信息

官方教程：《PYTORCH PROFILER WITH TENSORBOARD》

with torch.profiler.profile(
        schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
        on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/resnet18'),
        record_shapes=True,
        profile_memory=True,
        with_stack=True
) as prof:
    for step, batch_data in enumerate(train_loader):
        if step >= (1 + 1 + 3) * 2:
            break
        train(batch_data)
        prof.step()  # Call this at the end of each step to notify profiler of steps' boundary.