KinetoStepTracker

class torch.autograd.profiler.KinetoStepTracker

提供了一种全局递增步数（step count）的抽象。

以前，我们只能通过 pytorch profiler step() 在程序中的一个地方标记 step() 已经发生。现在我们将在 Optimizer 类中添加 step 钩子 https://github.com/pytorch/pytorch/issues/88446

• 这意味着已经对每次迭代调用 profiler.step() 的程序可能会导致步数被双重递增。

• 如果模型使用多个优化器，我们还可能导致步数被双重或更多地计数。

我们通过在调用 kineto 库的 step() 之前增加一个抽象层来解决这个问题。其思想是在一个字典中维护每个请求者的步数

{
   "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)

移除给定的请求者。

返回类型

bool

classmethod increment_step(requester)

递增请求者的步数。

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

返回类型

int

classmethod init_step_count(requester)

为给定的请求者初始化。