torch_npu.profiler._KinetoProfile

接口原型

1

torch_npu.profiler._KinetoProfile(activities=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None)

功能描述

提供PyTorch训练过程中的性能数据采集功能。

参数说明

• activities：CPU、NPU事件采集列表，Enum类型。可取值以及含义详见torch_npu.profiler.ProfilerActivity。
• record_shapes：算子的InputShapes和InputTypes，Bool类型。取值为：
  • True：开启。
  • False：关闭。

  默认值为False。

  开启torch_npu.profiler.ProfilerActivity.CPU时生效。

• profile_memory算子的内存占用情况，Bool类型。取值为：
  • True：开启。
  • False：关闭。

  默认值为False。

  

  已知在安装有glibc<2.34的环境上采集memory数据，可能触发glibc的一个已知Bug 19329，通过升级环境的glibc版本可解决此问题。

• with_stack：算子调用栈，Bool类型。包括框架层及CPU算子层的调用信息。取值为：
  • True：开启。
  • False：关闭。

  默认值为False。

  开启torch_npu.profiler.ProfilerActivity.CPU时生效。

• with_flops：算子浮点操作，Bool类型（该参数暂不支持解析性能数据）。取值为：
  • True：开启。
  • False：关闭。

  默认值为False。

  开启torch_npu.profiler.ProfilerActivity.CPU时生效。

• experimental_config：扩展参数，通过扩展配置性能分析工具常用的采集项。支持采集项和详细介绍请参见torch_npu.profiler._ExperimentalConfig。

支持的型号

• Atlas 训练系列产品
• Atlas A2 训练系列产品
• Atlas A3 训练系列产品

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

import torch
import torch_npu

...

prof = torch_npu.profiler._KinetoProfile(activities=None, record_shapes=False, profile_memory=False, with_stack=False, with_flops=False, with_modules=False, experimental_config=None)
for epoch in range(epochs):
        train_model_step()
        if epoch == 0:
                prof.start()
        if epoch == 1:
                prof.stop()
prof.export_chrome_trace("result_dir/trace.json")