feat: Improve layer naming

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -4,16 +4,18 @@
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set
 
 import numpy
+
+# @manual=//deeplearning/trt/python:py_tensorrt
+import tensorrt as trt
 import torch
 import torch.fx
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes.shape_prop import TensorMetadata
 from torch_tensorrt._Input import Input
+from torch_tensorrt.dynamo.conversion.converter_utils import get_node_name
 from torch_tensorrt.fx.observer import Observer
 from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
-# @manual=//deeplearning/trt/python:py_tensorrt
-import tensorrt as trt
 from packaging import version
 
 from .converter_registry import DYNAMO_CONVERTERS as CONVERTERS
@@ -232,7 +234,7 @@ def run(
         )
 
     def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
-        self._cur_node_name = str(n)
+        self._cur_node_name = get_node_name(n)
         self._cur_node = n
         # add "_itensor_to_tensor_meta"
         kwargs = dict(n.kwargs)

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -1,13 +1,40 @@
+import logging
+import re
 from typing import List
 
+import tensorrt as trt
 import torch
 from torch_tensorrt.fx.converters.converter_utils import (
     Frameworks,
     unified_dtype_converter,
 )
 from torch_tensorrt.fx.types import TRTDataType, TRTNetwork, TRTTensor
 
-import tensorrt as trt
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+def get_node_name(node: torch.fx.Node) -> str:
+    # nn_module_stack preserves the call stack of pytorch nn.modules
+    # The call stack contains a detailed name of the module
+    # which shows exactly where the module is located in the
+    # network architecture.
+    stack_item = node.meta.get("nn_module_stack", None)
+    # The current node is the last item in the stack
+    mod_stack = stack_item.popitem() if stack_item else ""
+    node_name = str(node)
+    if mod_stack:
+        mod_name = str(mod_stack[0]).replace("___", "/")
+        # Clean up the module name
+        mod_name = re.sub("^.*__self", "", mod_name)
+        mod_name = re.sub(r"_(\d+)$", r"/\g<1>", mod_name)
+        node_name = mod_name + "/" + node_name
+    else:
+        # Try an alternative way to get the module info
+        # like the node.meta['source_fn'] attr
+        pass
+
+    _LOGGER.debug(f"Node meta name {node_name}")
+    return node_name
 
 
 def dynamic_unsupported(node: torch.fx.Node) -> bool: