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参数的执行操作。

参数说明

• dir_name：采集的性能数据的输出目录，string类型。路径格式仅支持由字母、数字和下划线组成的字符串，不支持软链接。若配置tensorboard_trace_handler函数后未指定具体路径，性能数据默认落盘在当前目录。可选。
• worker_name：用于区分唯一的工作线程，string类型，默认为{hostname}_{pid}。路径格式仅支持由字母、数字和下划线组成的字符串，不支持软链接。可选。
• analyse_flag：性能数据自动解析开关，bool类型。可取值True（开启自动解析，默认值）、False（关闭自动解析，采集完后的性能数据可以使用离线解析）。可选。

• 离线解析详见《CANN 性能分析工具指南》中的“Ascend PyTorch Profiler性能数据离线解析”章节。

调用示例

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()