torch_npu.profiler.ProfilerActivity
接口原型
torch_npu.profiler.ProfilerActivity
功能描述
事件采集列表,Enum类型。用于赋值给torch_npu.profiler.profile的activities参数。
参数说明
- torch_npu.profiler.ProfilerActivity.CPU:框架侧数据采集的开关。
- torch_npu.profiler.ProfilerActivity.NPU:CANN软件栈及NPU数据采集的开关。
默认情况下两个开关同时开启。
支持的型号
- Atlas 训练系列产品
- Atlas A2 训练系列产品
调用示例
import torch
import torch_npu
...
with torch_npu.profiler.profile(
activities=[
torch_npu.profiler.ProfilerActivity.CPU,
torch_npu.profiler.ProfilerActivity.NPU
],
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result")
) as prof:
for step in range(steps):
train_one_step() # 训练函数
prof.step()
父主题: torch_npu.profiler