torch_npu.profiler.schedule

API接口

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

功能描述

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

参数说明

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

示例

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