KinetoStepTracker

class torch.autograd.profiler.KinetoStepTracker

提供了一个全局递增步数的抽象。

以前，我们只有一个地方可以通过 pytorch profiler step() 来标记程序中发生了 step()。现在，我们将在 Optimizer 类中添加 step 钩子 pytorch/pytorch#88446

• 这可能意味着已经每迭代调用 profiler.step() 的程序可能会重复计算步数。

• 如果模型使用多个优化器，也可能导致步数翻倍或更多。

我们通过在调用 step() 前添加一个抽象层来修复这个问题，该层会调用 kineto 库。思路是使用字典来维护每个请求者的步数。

{
   "ProfilerStep": 100,  # triggered by profiler step() call
   "Optimizer1Step": 100,   # Optimizer 1 or 2 are just examples, could be SGD, Adam etc
   "Optimizer2Step": 100,
}

要获取全局步数，只需取字典值中的最大值（100）。

如果其中一个计数器增加，最大值就会随之增加。

{
   "ProfilerStep": 100,
   "Optimizer1Step": 101,   # Optimizer1 got incremented first say
   "Optimizer2Step": 100,
}

此时全局步数为 101。我们仅在全局计数器增加时才调用 kineto step() 函数。

注意：目前请不要在 Optimizer 之外的模块中使用 KinetoStepTracker。这可能会导致不正确的步数增加。

classmethod current_step()

获取任何请求者的最新步数。

返回类型

int

classmethod erase_step_count(requester)

删除指定的请求者。

返回类型

布尔值

classmethod increment_step(requester)

递增指定请求者的步数。

此外，如果所有步数中的最大值增加，则触发 _kineto_step() 并返回全局步数。

返回类型

int

classmethod init_step_count(requester)

初始化指定请求者。