如何不使用上下文管理器调用torch.profiler.profile？

版本问题说明 #

Torch 1.9.0 确实移除了profile对象的start()/stop()方法，该版本的profile仅支持作为上下文管理器使用，官方部分教程的start/stop用法对应更高版本或旧版API，不适用于1.9.0。

无侵入式性能分析实现（无需修改训练脚本） #

你可以通过猴子补丁在训练开始前包装训练入口函数，将Profiler上下文管理器注入其中，无需改动训练脚本本身的代码：

假设训练脚本的核心训练逻辑入口是train_loop()函数，在训练前的初始化代码中执行以下操作：

from torch.profiler import profile, ProfilerActivity

# 保存原始训练函数引用
original_train = train_loop

def wrapped_train(*args, **kwargs):
    # 初始化Profiler，按需配置监控项
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True
    ) as prof:
        # 执行原始训练逻辑
        result = original_train(*args, **kwargs)
    
    # 输出性能分析结果，或导出到文件
    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
    prof.export_chrome_trace("training_profile.json")  # 可导入Chrome浏览器分析
    
    return result

# 替换原始训练函数为包装后的版本
train_loop = wrapped_train

当训练脚本调用train_loop()时，会自动触发Profiler的监控逻辑，全程无需修改训练脚本内部代码。

补充说明 #

• 如果训练逻辑没有统一的入口函数，也可以针对性包装模型的forward方法或优化器的step方法，实现局部性能监控；
• 1.9.0版本的profile上下文管理器支持schedule参数，可配置周期性采样（比如仅监控特定epoch），进一步适配你的分析需求。

内容的提问来源于stack exchange，提问作者tiberius