torch_npu.profiler.schedule

接口原型

torch_npu.profiler.schedule (wait, active, warmup = 0, repeat = 0, skip_first = 0)

功能描述

设置不同step的行为。用于构造torch_npu.profiler.profile的schedule参数。

参数说明

• wait：每次重复执行采集跳过的step轮数。必选。
• active：采集的step轮数。必选。
• warmup：预热的step轮数。默认值为0。建议设置1轮预热。可选。
• repeat：重复执行wait+warmup+active的次数。默认为值0。默认不执行该操作。可选。
• skip_first：采集前先跳过的step轮数。默认为值0。建议跳过前10轮。可选。

调用示例

with torch_npu.profiler.profile(
schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=1),
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()