ModulePathTracer is an FX tracer that--for each operation--also records the qualified name of the Module from which the operation originated.
| 13 | from typing import Any, Callable, Dict, Optional, Tuple |
| 14 | |
| 15 | class ModulePathTracer(torch.fx.Tracer): |
| 16 | """ |
| 17 | ModulePathTracer is an FX tracer that--for each operation--also records |
| 18 | the qualified name of the Module from which the operation originated. |
| 19 | """ |
| 20 | |
| 21 | # The current qualified name of the Module being traced. The top-level |
| 22 | # module is signified by empty string. This is updated when entering |
| 23 | # call_module and restored when exiting call_module |
| 24 | current_module_qualified_name : str = '' |
| 25 | # A map from FX Node to the qualname of the Module from which it |
| 26 | # originated. This is recorded by `create_proxy` when recording an |
| 27 | # operation |
| 28 | node_to_originating_module : Dict[torch.fx.Node, str] = {} |
| 29 | |
| 30 | def call_module(self, m: torch.nn.Module, forward: Callable[..., Any], |
| 31 | args : Tuple[Any, ...], kwargs : Dict[str, Any]) -> Any: |
| 32 | """ |
| 33 | Override of Tracer.call_module (see |
| 34 | https://pytorch.org/docs/stable/fx.html#torch.fx.Tracer.call_module). |
| 35 | |
| 36 | This override: |
| 37 | 1) Stores away the qualified name of the caller for restoration later |
| 38 | 2) Installs the qualified name of the caller in `current_module_qualified_name` |
| 39 | for retrieval by `create_proxy` |
| 40 | 3) Delegates into the normal Tracer.call_module method |
| 41 | 4) Restores the caller's qualified name into current_module_qualified_name |
| 42 | """ |
| 43 | old_qualname = self.current_module_qualified_name |
| 44 | try: |
| 45 | self.current_module_qualified_name = self.path_of_module(m) |
| 46 | return super().call_module(m, forward, args, kwargs) |
| 47 | finally: |
| 48 | self.current_module_qualified_name = old_qualname |
| 49 | |
| 50 | def create_proxy(self, kind: str, target: torch.fx.node.Target, args: Tuple[Any, ...], |
| 51 | kwargs: Dict[str, Any], name: Optional[str] = None, type_expr: Optional[Any] = None): |
| 52 | """ |
| 53 | Override of `Tracer.create_proxy`. This override intercepts the recording |
| 54 | of every operation and stores away the current traced module's qualified |
| 55 | name in `node_to_originating_module` |
| 56 | """ |
| 57 | proxy = super().create_proxy(kind, target, args, kwargs, name, type_expr) |
| 58 | self.node_to_originating_module[proxy.node] = self.current_module_qualified_name |
| 59 | return proxy |
| 60 | |
| 61 | |
| 62 | # Testing: let's see how this works on a torchvision ResNet18 model |