Fixed accuracy issue of fast refit

Author: cehongwang

py/torch_tensorrt/dynamo/_refit.py

@@ -117,10 +117,16 @@ def construct_refit_mapping_from_weight_name_map(
         torch_dtype = dtype.try_from(np_weight_type).to(torch.dtype)
         if engine_weight_name.split(" ")[-1] in ["SCALE", "SHIFT"]:
             # Batch Norm Layer
-            params = {}
+            params = {
+                "weight": 1.0,
+                "bias": 0.0,
+                "running_mean": 0.0,
+                "running_var": 1.0,
+            }
             for w in sd_weight_name:
-                params[w.split(".")[-1]] = state_dict[w]
-            scale = params["weight"] / torch.sqrt(params["running_var"] + 1e-7)
+                if w in state_dict:
+                    params[w.split(".")[-1]] = state_dict[w]
+            scale = params["weight"] / torch.sqrt(params["running_var"] + 1e-5)
             shift = params["bias"] - params["running_mean"] * scale
             # Set scale to scale or shift to shift
             engine_weight_map[engine_weight_name] = eval(
@@ -171,6 +177,11 @@ def _refit_single_trt_engine_with_gm(
         mapping = construct_refit_mapping_from_weight_name_map(
             weight_name_map, new_gm.state_dict()
         )
+
+        # Debug Use
+        # correct = construct_refit_mapping(new_gm, input_list, settings)
+        # {k: np.allclose(correct[k][0], mapping[k][0].cpu().numpy(), 1e-2, 1e-2) for k in mapping if k in correct}
+
         for layer_name in weight_list:
             if layer_name not in mapping:
                 logger.warning(f"{layer_name} is not found in weight mapping.")