torch_npu.profiler.profile
API接口
torch_npu.profiler.profile(activities=None, schedule=None, on_trace_ready=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None, use_cuda=None)
功能描述
提供对训练过程数据的profiling功能。
参数说明
- schedule:设置不同step的行为,Callable类型。由schedule类控制。
- on_trace_ready:采集结束时自动执行操作,Callable类型。当前仅支持执行tensorboard_trace_handler函数的操作,默认不执行任何操作。
- record_shapes:算子的InputShapes和InputTypes,Bool类型。取值为:
- profile_memory算子的内存占用情况,Bool类型。取值为:
- True:开启。
- False:关闭。默认值。
- with_stack:算子调用栈,Bool类型。取值为:
- with_flops:算子浮点操作,Bool类型(该参数暂不支持解析性能数据)。取值为:
- with_modules:with_stack时modules分层信息,Bool类型(该参数暂不支持解析性能数据)。取值为:
- experimental_config:扩展参数,通过扩展配置性能分析工具常用的采集项。支持采集项和详细介绍请参见torch_npu.profiler._ExperimentalConfig
示例
experimental_config = torch_npu.profiler._ExperimentalConfig( aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False ) with torch_npu.profiler.profile( activities=[ torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU ], schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=2, repeat=2, skip_first=10), on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result"), record_shapes=True, profile_memory=True, with_stack=True, with_flops=False, with_modules=False, experimental_config=experimental_config) as prof: for step in range(steps): train_one_step(step, steps, train_loader, model, optimizer, criterion) prof.step()
父主题: torch_npu.profiler