KinetoStepTracker

class torch.autograd.profiler.KinetoStepTracker

提供了一个抽象方法来全局递增步数。

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

• 这意味着已经每迭代调用 profiler.step() 的程序可能会导致步骤计数被双倍增加。

• 如果一个模型使用多个优化器，我们也可以有双倍或更多的步骤计数。

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

{
   "ProfilerStep": 100,  # 由 profiler step() 调用触发
   "Optimizer1Step": 100,   # Optimizer 1 或 2 只是示例，可能是 SGD、Adam 等
   "Optimizer2Step": 100,
}

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

如果其中一个计数增加，最大值将会上升。

{
   "ProfilerStep": 100,
   "Optimizer1Step": 101,   # 优化器1首先增加
   "Optimizer2Step": 100,
}

然后全局步数为101 我们只在全局计数增加时调用kineto step()函数。

注意：目前请不要在优化器以外的模块中使用KinetoStepTracker。结果可能会导致步数不正确地增加。

classmethod current_step()

获取任何请求者的最新步骤

Return type

int

classmethod erase_step_count(requester)

移除指定的请求者。

Return type

bool

classmethod increment_step(requester)

增加请求者的步骤计数。

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

Return type

int

classmethod init_step_count(requester)

为指定的请求者进行初始化。