Trained weight loading for unit tests

Author: ign-navaneethk

models/experimental/panoptic_deeplab/common.py

@@ -11,6 +11,11 @@
 from models.experimental.panoptic_deeplab.reference.resnet52_backbone import ResNet52BackBone as TorchBackbone
 from models.experimental.panoptic_deeplab.reference.resnet52_stem import DeepLabStem
 from torchvision.models.resnet import Bottleneck
+from models.experimental.panoptic_deeplab.reference.aspp import ASPPModel
+from models.experimental.panoptic_deeplab.reference.decoder import DecoderModel
+from models.experimental.panoptic_deeplab.reference.res_block import ResModel
+from models.experimental.panoptic_deeplab.reference.head import HeadModel
+from models.experimental.panoptic_deeplab.reference.panoptic_deeplab import TorchPanopticDeepLab
 
 
 def map_single_key(checkpoint_key):
@@ -259,17 +264,18 @@ def map_single_key(checkpoint_key):
     return key
 
 
-def load_partial_state(torch_model: torch.nn.Module, state_dict, layer_prefix: str = ""):
+def load_partial_state(torch_model: torch.nn.Module, state_dict, layer_name: str = ""):
     partial_state_dict = {}
+    layer_prefix = layer_name + "."
     for k, v in state_dict.items():
         if k.startswith(layer_prefix):
             partial_state_dict[k[len(layer_prefix) :]] = v
     torch_model.load_state_dict(partial_state_dict, strict=True)
     logger.info(f"Successfully loaded all mapped weights with strict=True")
-    return torch_model.eval()
+    return torch_model
 
 
-def load_torch_model_state(model: torch.nn.Module = None, layer_name: str = "", model_location_generator=None):
+def load_torch_model_state(torch_model: torch.nn.Module = None, layer_name: str = "", model_location_generator=None):
     if model_location_generator == None or "TT_GH_CI_INFRA" not in os.environ:
         model_path = "models"
     else:
@@ -312,16 +318,24 @@ def load_torch_model_state(model: torch.nn.Module = None, layer_name: str = "",
     for checkpoint_key, model_key in key_mapping.items():
         mapped_state_dict[model_key] = state_dict[checkpoint_key]
 
-    if model is None:
-        return mapped_state_dict
-    elif isinstance(model, TorchBackbone):
-        layer_prefix = "backbone."
-        return load_partial_state(model, mapped_state_dict, layer_prefix)
-    elif isinstance(model, DeepLabStem):
-        layer_prefix = "backbone.stem."
-        return load_partial_state(model, mapped_state_dict, layer_prefix)
-    elif isinstance(model, Bottleneck):
-        layer_prefix = "backbone." + layer_name + "."
-        return load_partial_state(model, mapped_state_dict, layer_prefix)
+    if isinstance(
+        torch_model,
+        (
+            DeepLabStem,
+            Bottleneck,
+            TorchBackbone,
+            ASPPModel,
+            ResModel,
+            HeadModel,
+            DecoderModel,
+        ),
+    ):
+        torch_model = load_partial_state(torch_model, mapped_state_dict, layer_name)
+    elif isinstance(torch_model, TorchPanopticDeepLab):
+        del mapped_state_dict["pixel_mean"]
+        del mapped_state_dict["pixel_std"]
+        torch_model.load_state_dict(mapped_state_dict, strict=True)
     else:
         raise NotImplementedError("Unknown torch model. Weight loading not implemented")
+
+    return torch_model.eval()

models/experimental/panoptic_deeplab/reference/decoder.py

@@ -22,7 +22,7 @@ def __init__(self, name) -> None:
         super().__init__()
         self.name = name
         self.aspp = ASPPModel()
-        if name == "semantics_head":
+        if name == "semantic_decoder":
             self.res3 = ResModel(512, 320, 256)
             self.res2 = ResModel(256, 288, 256)
             self.head_1 = HeadModel(256, 256, 19)
@@ -48,7 +48,7 @@ def forward(self, x: Tensor, res3: Tensor, res2: Tensor) -> Tuple[Tensor, Tensor
         out_ = self.res2(out, res2)
         out = self.head_1(out_)
 
-        if self.name == "instance_head":
+        if self.name == "instance_decoder":
             out_2 = self.head_2(out_)
         else:
             out_2 = None

models/experimental/panoptic_deeplab/reference/head.py

@@ -22,20 +22,24 @@ def __init__(self, in_channels, intermediate_channels, out_channels) -> None:
 
         if out_channels == 1:  # instance center head
             self.conv1 = nn.Sequential(
-                nn.Conv2d(in_channels, in_channels, 3, 1, 1, 1), nn.BatchNorm2d(in_channels), nn.ReLU()
+                nn.Conv2d(in_channels, in_channels, 3, 1, 1, 1, bias=False), nn.BatchNorm2d(in_channels), nn.ReLU()
             )
 
             self.conv2 = nn.Sequential(
-                nn.Conv2d(in_channels, intermediate_channels, 3, 1, 1, 1),
+                nn.Conv2d(in_channels, intermediate_channels, 3, 1, 1, 1, bias=False),
                 nn.BatchNorm2d(intermediate_channels),
                 nn.ReLU(),
             )
         else:  # instance offset head and semantics head
             self.conv1 = nn.Sequential(
-                nn.Conv2d(in_channels, in_channels, 5, 1, 2, 1, in_channels), nn.BatchNorm2d(in_channels), nn.ReLU()
+                nn.Conv2d(in_channels, in_channels, 5, 1, 2, 1, in_channels, bias=False),
+                nn.BatchNorm2d(in_channels),
+                nn.ReLU(),
             )
             self.conv2 = nn.Sequential(
-                nn.Conv2d(in_channels, intermediate_channels, 1, 1), nn.BatchNorm2d(intermediate_channels), nn.ReLU()
+                nn.Conv2d(in_channels, intermediate_channels, 1, 1, bias=False),
+                nn.BatchNorm2d(intermediate_channels),
+                nn.ReLU(),
             )
         self.conv3 = nn.Sequential(nn.Conv2d(intermediate_channels, out_channels, 1, 1))
 

models/experimental/panoptic_deeplab/reference/panoptic_deeplab.py

@@ -20,17 +20,24 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        # self.pixel_std = nn.Parameter(torch.randn((3, 1, 1)))
+        # self.pixel_mean = nn.Parameter(torch.randn((3, 1, 1)))
+        # self.register_buffer("pixel_mean", torch.randn(3).view(-1, 1, 1), False)
+        # self.register_buffer("pixel_std", torch.randn(3).view(-1, 1, 1), False)
+        # self.register_buffer("adsaf", torch.randn(3).view(-1, 1, 1), False)
+        # self.register_buffer("yurfdgdf", torch.randn(3).view(-1, 1, 1), False)
+
         # Backbone
         self.backbone = ResNet52BackBone()
 
         # Semantic segmentation decoder
         self.semantic_decoder = DecoderModel(
-            name="Semantics_head",
+            name="semantic_decoder",
         )
 
         # Instance segmentation decoders
         self.instance_decoder = DecoderModel(
-            name="instance_head",
+            name="instance_decoder",
         )
 
     def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

models/experimental/panoptic_deeplab/tests/test_aspp.py

@@ -9,10 +9,9 @@
 
 from tests.ttnn.utils_for_testing import check_with_pcc
 from models.experimental.panoptic_deeplab.tt.custom_preprocessing import create_custom_mesh_preprocessor
-from models.experimental.panoptic_deeplab.reference.aspp import (
-    ASPPModel,
-)
+from models.experimental.panoptic_deeplab.reference.aspp import ASPPModel
 from models.experimental.panoptic_deeplab.tt.aspp import TTASPP
+from models.experimental.panoptic_deeplab.common import load_torch_model_state
 
 
 class AsppTestInfra:
@@ -24,6 +23,7 @@ def __init__(
         height,
         width,
         model_config,
+        name,
     ):
         super().__init__()
         if not hasattr(self, "_model_initialized"):
@@ -37,19 +37,20 @@ def __init__(
         self.num_devices = device.get_num_devices()
         self.batch_size = batch_size
         self.inputs_mesh_mapper, self.weights_mesh_mapper, self.output_mesh_composer = self.get_mesh_mappers(device)
+        self.name = name
 
         # torch model
-        torch_model = ASPPModel().eval()
-        self.torch_input_tensor = torch.randn((batch_size, input_channels, height, width), dtype=torch.float16)
+        torch_model = ASPPModel()
+        torch_model = load_torch_model_state(torch_model, name)
+
         parameters = preprocess_model_parameters(
             initialize_model=lambda: torch_model,
             custom_preprocessor=create_custom_mesh_preprocessor(self.weights_mesh_mapper),
             device=None,
         )
-        torch_model.to(torch.bfloat16)
-        self.torch_input_tensor = self.torch_input_tensor.to(torch.bfloat16)
 
         # golden
+        self.torch_input_tensor = torch.randn((batch_size, input_channels, height, width), dtype=torch.float)
         self.torch_output_tensor = torch_model(self.torch_input_tensor)
 
         # ttnn
@@ -67,7 +68,6 @@ def __init__(
         # run and validate
         self.run()
         self.validate()
-        ttnn.deallocate(self.output_tensor)
 
     def get_mesh_mappers(self, device):
         if device.get_num_devices() != 1:
@@ -85,23 +85,30 @@ def run(self):
         return self.output_tensor
 
     def validate(self, output_tensor=None):
-        """Validate outputs"""
+        tt_output_tensor = self.output_tensor if output_tensor is None else output_tensor
+        tt_output_tensor_torch = ttnn.to_torch(
+            tt_output_tensor, device=self.device, mesh_composer=self.output_mesh_composer
+        )
+
+        # Deallocate output tesnors
+        ttnn.deallocate(tt_output_tensor)
 
-        output_tensor = self.output_tensor if output_tensor is None else output_tensor
-        output_tensor = ttnn.to_torch(output_tensor, device=self.device, mesh_composer=self.output_mesh_composer)
         expected_shape = self.torch_output_tensor.shape
-        output_tensor = torch.reshape(
-            output_tensor, (expected_shape[0], expected_shape[2], expected_shape[3], expected_shape[1])
+        tt_output_tensor_torch = torch.reshape(
+            tt_output_tensor_torch, (expected_shape[0], expected_shape[2], expected_shape[3], expected_shape[1])
         )
-        output_tensor = torch.permute(output_tensor, (0, 3, 1, 2))
-        batch_size = self.batch_size
+        tt_output_tensor_torch = torch.permute(tt_output_tensor_torch, (0, 3, 1, 2))
+
+        batch_size = tt_output_tensor_torch.shape[0]
 
-        valid_pcc = 0.97
-        self.pcc_passed, self.pcc_message = check_with_pcc(self.torch_output_tensor, output_tensor, pcc=valid_pcc)
+        valid_pcc = 0.99
+        self.pcc_passed, self.pcc_message = check_with_pcc(
+            self.torch_output_tensor, tt_output_tensor_torch, pcc=valid_pcc
+        )
 
         assert self.pcc_passed, logger.error(f"PCC check failed: {self.pcc_message}")
         logger.info(
-            f"Modular Panoptic DeepLab ASPP - batch_size={batch_size}, act_dtype={model_config['ACTIVATIONS_DTYPE']}, weight_dtype={model_config['WEIGHTS_DTYPE']}, math_fidelity={model_config['MATH_FIDELITY']}, PCC={self.pcc_message}"
+            f"Modular Panoptic DeepLab ASPP layer:{self.name} - batch_size={batch_size}, act_dtype={model_config['ACTIVATIONS_DTYPE']}, weight_dtype={model_config['WEIGHTS_DTYPE']}, math_fidelity={model_config['MATH_FIDELITY']}, PCC={self.pcc_message}"
         )
 
         return self.pcc_passed, self.pcc_message
@@ -121,12 +128,14 @@ def validate(self, output_tensor=None):
         (1, 2048, 32, 64),
     ],
 )
-def test_aspp(device, batch_size, input_channels, height, width):
+@pytest.mark.parametrize("name", ["semantic_decoder.aspp", "instance_decoder.aspp"])
+def test_aspp(device, batch_size, input_channels, height, width, name):
     AsppTestInfra(
         device,
         batch_size,
         input_channels,
         height,
         width,
         model_config,
+        name,
     )

models/experimental/panoptic_deeplab/tests/test_decoder.py

@@ -13,9 +13,8 @@
     decoder_layer_optimisations,
 )
 from models.experimental.panoptic_deeplab.tt.custom_preprocessing import create_custom_mesh_preprocessor
-from models.experimental.panoptic_deeplab.reference.decoder import (
-    DecoderModel,
-)
+from models.experimental.panoptic_deeplab.reference.decoder import DecoderModel
+from models.experimental.panoptic_deeplab.common import load_torch_model_state
 
 
 class DecoderTestInfra:
@@ -56,28 +55,26 @@ def __init__(
 
         # Create input tensors
         self.torch_input_tensor = torch.randn(
-            (self.batch_size, self.in_channels, self.height, self.width), dtype=torch.float32
+            (self.batch_size, self.in_channels, self.height, self.width), dtype=torch.float
         )
 
         # Create res3 and res2 feature maps with appropriate dimensions
-        self.torch_res3_tensor = torch.randn(
-            (self.batch_size, 512, self.height * 2, self.width * 2), dtype=torch.float32
-        )
+        self.torch_res3_tensor = torch.randn((self.batch_size, 512, self.height * 2, self.width * 2), dtype=torch.float)
 
         self.torch_res2_tensor = torch.randn(
-            (self.batch_size, upsample_channels, self.height * 4, self.width * 4), dtype=torch.float32
+            (self.batch_size, upsample_channels, self.height * 4, self.width * 4), dtype=torch.float
         )
 
         # torch model
-        torch_model = DecoderModel(self.name).eval()
+        torch_model = DecoderModel(self.name)
+        torch_model = load_torch_model_state(torch_model, name)
 
         parameters = preprocess_model_parameters(
             initialize_model=lambda: torch_model,
             custom_preprocessor=create_custom_mesh_preprocessor(self.weights_mesh_mapper),
             device=None,
         )
 
-        parameters.conv_args = {}
         # For ASPP
         aspp_args = infer_ttnn_module_args(
             model=torch_model.aspp, run_model=lambda model: model(self.torch_input_tensor), device=None
@@ -116,12 +113,6 @@ def __init__(
             if hasattr(parameters, "head_2"):
                 parameters.head_2.conv_args = head_2_args
 
-        # Convert to bfloat16
-        torch_model.to(torch.bfloat16)
-        self.torch_input_tensor = self.torch_input_tensor.to(torch.bfloat16)
-        self.torch_res3_tensor = self.torch_res3_tensor.to(torch.bfloat16)
-        self.torch_res2_tensor = self.torch_res2_tensor.to(torch.bfloat16)
-
         # Get torch output with all three inputs
         self.torch_output_tensor, self.torch_output_tensor_2 = torch_model(
             self.torch_input_tensor, self.torch_res3_tensor, self.torch_res2_tensor
@@ -189,7 +180,7 @@ def validate(self, output_tensor=None):
 
         batch_size = output_tensor.shape[0]
 
-        valid_pcc = 0.97
+        valid_pcc = 0.99
         self.pcc_passed, self.pcc_message = check_with_pcc(self.torch_output_tensor, output_tensor, pcc=valid_pcc)
         assert self.pcc_passed, logger.error(f"PCC check failed: {self.pcc_message}")
         logger.info(
@@ -207,7 +198,7 @@ def validate(self, output_tensor=None):
 
             batch_size = output_tensor.shape[0]
 
-            valid_pcc = 0.96
+            valid_pcc = 0.99
             self.pcc_passed, self.pcc_message = check_with_pcc(self.torch_output_tensor_2, output_tensor, pcc=valid_pcc)
             assert self.pcc_passed, logger.error(f"PCC check failed: {self.pcc_message}")
             logger.info(
@@ -228,8 +219,8 @@ def validate(self, output_tensor=None):
 @pytest.mark.parametrize(
     "batch_size, in_channels, res3_intermediate_channels, res2_intermediate_channels, out_channels, upsample_channels, height, width, name",
     [
-        (1, 2048, 320, 288, (19,), 256, 32, 64, "semantics_head"),  # semantic head
-        (1, 2048, 320, 160, (2, 1), 256, 32, 64, "instance_head"),  # instance offset head
+        (1, 2048, 320, 288, (19,), 256, 32, 64, "semantic_decoder"),  # semantic head
+        (1, 2048, 320, 160, (2, 1), 256, 32, 64, "instance_decoder"),  # instance offset head
     ],
 )
 def test_decoder(