torch_npu.profiler.tensorboard_trace_handler

接口原型

torch_npu.profiler.tensorboard_trace_handler(dir_name=None, worker_name=None, analyse_flag=True)

功能描述

将采集到的性能数据导出为TensorBoard工具支持的格式。作为torch_npu.profiler.profile on_trace_ready参数的执行操作。

参数说明

调用示例

import torch
import torch_npu

...

with torch_npu.profiler.profile(
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result")
    ) as prof:
            for step in range(steps):
                train_one_step(step, steps, train_loader, model, optimizer, criterion)
                prof.step()