torch.utils.module_tracker

此实用程序可用于跟踪层次结构中的当前位置。 它可以在其他跟踪工具中使用，以便能够轻松地将测量的数量与用户友好的名称相关联。这在今天的 FlopCounterMode 中特别使用。

torch.utils.module_tracker 类。ModuleTracker

ModuleTracker是跟踪 NN.执行期间的模块层次结构 以便其他系统可以查询当前正在执行的 Module （或其向后执行的 执行）。

您可以访问此上下文管理器上的属性，以获取所有 当前通过其 fqn（完全限定名称，也用作 state_dict）。 您可以访问该属性以了解您当前是否在 backward 中运行。parentsis_bw

请注意，它永远不会为空，并且始终包含 “Global” 键。标志 将在 forward 之后保持，直到执行另一个 Module。如果你需要它 更准确地说，请提交一个请求此的 issue。将映射从 fqn 添加到模块实例 是可能的，但尚未完成，如果需要，请提交一个请求此 Issue。parentsis_bwTrue

用法示例

mod = torch.nn.Linear(2, 2)

with ModuleTracker() as tracker:
    # Access anything during the forward pass
    def my_linear(m1, m2, bias):
        print(f"Current modules: {tracker.parents}")
        return torch.mm(m1, m2.t()) + bias
    torch.nn.functional.linear = my_linear

    mod(torch.rand(2, 2))