Added comments and fixed some issue

Author: cehongwang

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -323,6 +323,13 @@ def _construct_trt_network_def(self) -> None:
         )
 
     def _save_weight_mapping(self) -> None:
+        """
+        Construct the weight name mapping from engine weight name to state_dict weight name.
+        Cache the weight name for future refitting usecases.
+        Two-stage weight name tracing:
+        1. Name transformation from engine weight name to state_dict weight name
+        2. Value mapping that, for each weight in INetworkDefinition search for identical weight in state_dict
+        """
 
         def find_weight(
             weight_name: str, np_map: dict[str, Any], sd: dict[str, Any]
@@ -386,7 +393,7 @@ def check_weight_equal(
             )
         }
         """
-
+        # Stage 1: Name mapping
         sd = self.module.state_dict()
         weight_name_map: dict[str, Any] = {}
         np_map = {}
@@ -413,6 +420,7 @@ def check_weight_equal(
                         [i for i in sd_weight_name_list[:-1] if i]
                     )
                     suffix = sd_weight_name_list[-1]
+                    # Retrieve each weight name(s) in state_dict
                     if layer_type == "CONSTANT":
                         if "embedding" in suffix:
                             sd_weight_name = f"{sd_weight_name}.{torch_attr[0]}"
@@ -430,7 +438,7 @@ def check_weight_equal(
                     weight_name_map[engine_weight_name] = sd_weight_name
                     np_map[engine_weight_name] = weight
 
-        # Value mapping
+        # Stage 2: Value mapping
         for engine_weight_name, sd_weight_name in weight_name_map.items():
             if "SCALE" in engine_weight_name:
                 # There is no direct connection in batch_norm layer. So skip it
@@ -448,7 +456,6 @@ def check_weight_equal(
             ]
 
         self.weight_name_map = weight_name_map
-        # check = {k:(weight_name_map[k], np_map[k]) for k, v in np_map.items()}
 
     def run(
         self,

tests/py/dynamo/models/test_model_refit.py

@@ -127,7 +127,7 @@ def test_fast_refit_one_engine():
 
 
 @pytest.mark.unit
-def test_fast_refit_one_engin_no_map():
+def test_fast_refit_one_engine_no_map():
 
     model = models.resnet18(pretrained=False).eval().to("cuda")
     model2 = models.resnet18(pretrained=True).eval().to("cuda")
@@ -174,7 +174,7 @@ def test_fast_refit_one_engin_no_map():
 
 
 @pytest.mark.unit
-def test_fast_refit_one_engin_wrong_map():
+def test_fast_refit_one_engine_wrong_map():
 
     model = models.resnet18(pretrained=False).eval().to("cuda")
     model2 = models.resnet18(pretrained=True).eval().to("cuda")