torch_npu.profiler.tensorboard_trace_handler

接口原型

torch_npu.profiler.tensorboard_trace_handler(dir_name, worker_name=None, use_gzip=False)

功能描述

将采集到的性能数据导出为TensorBoard工具支持的格式。用于构造torch_npu.profiler.profile的on_trace_ready参数。

参数说明

• dir_name：采集的性能数据的输出目录。必选。
• worker_name：用于区分唯一的工作线程，默认为{hostname}_{pid}。可选。
• use_gzip：暂未使用。

调用示例

with torch_npu.profiler.profile(
schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=1),
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result"),
    with_stack=True) as prof:
            for step in range(steps):
                train_one_step(step, steps, train_loader, model, optimizer, criterion)
                prof.step()