torch_npu.profiler.tensorboard_trace_handler

API接口

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=1, warmup=1, active=2, repeat=2, skip_first=10),
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()