torch.utils.module_tracker¶
此实用程序可用于跟踪 torch.nn.Module
层级结构中的当前位置。它可以在其他跟踪工具中使用,以便能够轻松地将测量的量与用户友好的名称相关联。这在 FlopCounterMode 中尤其有用。
- class torch.utils.module_tracker.ModuleTracker[source]¶
ModuleTracker
是一个上下文管理器,它在执行期间跟踪 nn.Module 层级结构,以便其他系统可以查询当前正在执行的模块(或其反向传播正在执行)。您可以访问此上下文管理器的
parents
属性,以获取通过其 fqn(完全限定名,也用作 state_dict 中的键)当前正在执行的所有模块集。您可以访问is_bw
属性来了解您当前是否在反向传播中运行。请注意,
parents
从不为空,并且始终包含“全局”键。在执行另一个模块之前,is_bw
标志将在前向传播后保持True
。如果您需要它更准确,请提交一个问题请求此功能。添加一个从 fqn 到模块实例的映射是可能的,但尚未完成,如果您需要它,请提交一个问题请求此功能。示例用法
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))