refactored full net test

Author: ign-febin

models/experimental/panoptic_deeplab/tests/test_panoptic_deeplab.py

@@ -5,252 +5,175 @@
 import torch
 from loguru import logger
 import ttnn
-from ttnn.model_preprocessing import preprocess_model_parameters
+from ttnn.model_preprocessing import (
+    preprocess_model_parameters,
+    infer_ttnn_module_args,
+)
 from tests.ttnn.utils_for_testing import check_with_pcc
 
 from models.experimental.panoptic_deeplab.reference.panoptic_deeplab import TorchPanopticDeepLab
 from models.experimental.panoptic_deeplab.tt.panoptic_deeplab import TTPanopticDeepLab
 from models.experimental.panoptic_deeplab.tt.custom_preprocessing import create_custom_mesh_preprocessor
-from ttnn.model_preprocessing import infer_ttnn_module_args, preprocess_model_parameters
-from models.experimental.panoptic_deeplab.common import load_torch_model_state
 
 
 class PanopticDeepLabTestInfra:
-    def __init__(
-        self,
-        device,
-        batch_size,
-        in_channels,
-        height,
-        width,
-        model_config,
-    ):
+    _seeded = False
+    _PCC_THRESH = 0.97
+
+    def __init__(self, device, batch_size, in_channels, height, width, model_config):
         super().__init__()
-        if not hasattr(self, "_model_initialized"):
-            torch.manual_seed(42)
-            self._model_initialized = True
-            torch.cuda.manual_seed_all(42)
-            torch.backends.cudnn.deterministic = True
+        self._maybe_seed()
 
-        self.pcc_passed = False
-        self.pcc_message = "call validate()?"
+        # Core state
         self.device = device
+        self.model_config = model_config
         self.num_devices = device.get_num_devices()
-        self.batch_size = batch_size
-        self.in_channels = in_channels
-        self.height = height
-        self.width = width
+        self.batch_size, self.in_channels, self.height, self.width = (
+            batch_size,
+            in_channels,
+            height,
+            width,
+        )
         self.inputs_mesh_mapper, self.weights_mesh_mapper, self.output_mesh_composer = self.get_mesh_mappers(device)
 
-        # Initialize torch model
-        torch_model = TorchPanopticDeepLab()
-        torch_model = load_torch_model_state(torch_model, "panoptic_deeplab")
-
-        # Create input tensor
+        # Torch reference model + inputs
+        torch_model = TorchPanopticDeepLab().eval()
         input_shape = (batch_size * self.num_devices, in_channels, height, width)
         self.torch_input_tensor = torch.rand(input_shape, dtype=torch.float)
 
-        # Preprocess model parameters
+        # Preprocess TTNN parameters
         parameters = preprocess_model_parameters(
             initialize_model=lambda: torch_model,
             custom_preprocessor=create_custom_mesh_preprocessor(self.weights_mesh_mapper),
             device=None,
         )
 
-        parameters.conv_args = {}
-        input_tensor = torch.randn(1, 2048, 32, 64)
-        res3_tensor = torch.randn(1, 512, 64, 128)
-        res2_tensor = torch.randn(1, 256, 128, 256)
+        # Populate conv_args for decoders via one small warm-up pass
+        self._populate_all_decoders(torch_model, parameters)
 
-        # For semantic decoder
-        if hasattr(parameters, "semantic_decoder"):
-            # ASPP
-            aspp_args = infer_ttnn_module_args(
-                model=torch_model.semantic_decoder.aspp, run_model=lambda model: model(input_tensor), device=None
-            )
-            if hasattr(parameters.semantic_decoder, "aspp"):
-                parameters.semantic_decoder.aspp.conv_args = aspp_args
-
-            # Res3
-            aspp_out = torch_model.semantic_decoder.aspp(input_tensor)
-            res3_args = infer_ttnn_module_args(
-                model=torch_model.semantic_decoder.res3,
-                run_model=lambda model: model(aspp_out, res3_tensor),
-                device=None,
-            )
-            if hasattr(parameters.semantic_decoder, "res3"):
-                parameters.semantic_decoder.res3.conv_args = res3_args
-
-            # Res2
-            res3_out = torch_model.semantic_decoder.res3(aspp_out, res3_tensor)
-            res2_args = infer_ttnn_module_args(
-                model=torch_model.semantic_decoder.res2,
-                run_model=lambda model: model(res3_out, res2_tensor),
-                device=None,
-            )
-            if hasattr(parameters.semantic_decoder, "res2"):
-                parameters.semantic_decoder.res2.conv_args = res2_args
-
-            # Head
-            res2_out = torch_model.semantic_decoder.res2(res3_out, res2_tensor)
-            head_args = infer_ttnn_module_args(
-                model=torch_model.semantic_decoder.head_1, run_model=lambda model: model(res2_out), device=None
-            )
-            if hasattr(parameters.semantic_decoder, "head_1"):
-                parameters.semantic_decoder.head_1.conv_args = head_args
-
-        # For instance decoder
-        if hasattr(parameters, "instance_decoder"):
-            # ASPP
-            aspp_args = infer_ttnn_module_args(
-                model=torch_model.instance_decoder.aspp, run_model=lambda model: model(input_tensor), device=None
-            )
-            if hasattr(parameters.instance_decoder, "aspp"):
-                parameters.instance_decoder.aspp.conv_args = aspp_args
-
-            # Res3
-            aspp_out = torch_model.instance_decoder.aspp(input_tensor)
-            res3_args = infer_ttnn_module_args(
-                model=torch_model.instance_decoder.res3,
-                run_model=lambda model: model(aspp_out, res3_tensor),
-                device=None,
-            )
-            if hasattr(parameters.instance_decoder, "res3"):
-                parameters.instance_decoder.res3.conv_args = res3_args
-
-            # Res2
-            res3_out = torch_model.instance_decoder.res3(aspp_out, res3_tensor)
-            res2_args = infer_ttnn_module_args(
-                model=torch_model.instance_decoder.res2,
-                run_model=lambda model: model(res3_out, res2_tensor),
-                device=None,
-            )
-            if hasattr(parameters.instance_decoder, "res2"):
-                parameters.instance_decoder.res2.conv_args = res2_args
-
-            # Head
-            res2_out = torch_model.instance_decoder.res2(res3_out, res2_tensor)
-            head_args_1 = infer_ttnn_module_args(
-                model=torch_model.instance_decoder.head_1, run_model=lambda model: model(res2_out), device=None
-            )
-            head_args_2 = infer_ttnn_module_args(
-                model=torch_model.instance_decoder.head_2, run_model=lambda model: model(res2_out), device=None
-            )
-            if hasattr(parameters.instance_decoder, "head_1"):
-                parameters.instance_decoder.head_1.conv_args = head_args_1
-            if hasattr(parameters.instance_decoder, "head_2"):
-                parameters.instance_decoder.head_2.conv_args = head_args_2
-
-        # Run torch model with bfloat16
+        # Run Torch once (fp32) → then bf16 for parity with TTNN
         logger.info("Running PyTorch model...")
         self.torch_output_tensor, self.torch_output_tensor_2, self.torch_output_tensor_3 = torch_model(
             self.torch_input_tensor
         )
 
-        # Convert input to TTNN format (NHWC)
+        # Convert input to TTNN NHWC host tensor
         logger.info("Converting input to TTNN format...")
         tt_host_tensor = ttnn.from_torch(
             self.torch_input_tensor.permute(0, 2, 3, 1),
             dtype=ttnn.bfloat16,
             mesh_mapper=self.inputs_mesh_mapper,
         )
 
-        # Initialize TTNN model
+        # TTNN model
         logger.info("Initializing TTNN model...")
-        print("Initializing TTNN model...")
-        self.ttnn_model = TTPanopticDeepLab(
-            parameters=parameters,
-            model_config=model_config,
-        )
+        self.ttnn_model = TTPanopticDeepLab(parameters=parameters, model_config=model_config)
 
-        logger.info("Running first TTNN model pass (JIT configuration)...")
-        # first run configures convs JIT
-        self.input_tensor = ttnn.to_device(tt_host_tensor, device)
-        self.run()
-        self.validate()
+        # First run configures JIT, second run is optimized
+        for phase in ("JIT configuration", "optimized"):
+            logger.info(f"Running TTNN model pass ({phase})...")
+            self.input_tensor = ttnn.to_device(tt_host_tensor, device)
+            self.run()
+            self.validate()
 
-        logger.info("Running optimized TTNN model pass...")
-        # Optimized run
-        self.input_tensor = ttnn.to_device(tt_host_tensor, device)
-        self.run()
-        self.validate()
+    # --------------------------- Setup & helpers ---------------------------
+
+    @classmethod
+    def _maybe_seed(cls):
+        if not cls._seeded:
+            torch.manual_seed(42)
+            torch.cuda.manual_seed_all(42)
+            torch.backends.cudnn.deterministic = True
+            cls._seeded = True
 
     def get_mesh_mappers(self, device):
         if device.get_num_devices() != 1:
-            inputs_mesh_mapper = ttnn.ShardTensorToMesh(device, dim=0)
-            weights_mesh_mapper = None
-            output_mesh_composer = ttnn.ConcatMeshToTensor(device, dim=0)
-        else:
-            inputs_mesh_mapper = None
-            weights_mesh_mapper = None
-            output_mesh_composer = None
-        return inputs_mesh_mapper, weights_mesh_mapper, output_mesh_composer
+            return (
+                ttnn.ShardTensorToMesh(device, dim=0),  # inputs
+                None,  # weights
+                ttnn.ConcatMeshToTensor(device, dim=0),  # outputs
+            )
+        return None, None, None
 
-    def run(self):
-        self.output_tensor, self.output_tensor_2, self.output_tensor_3 = self.ttnn_model(self.input_tensor, self.device)
-        return self.output_tensor, self.output_tensor_2, self.output_tensor_3
+    @staticmethod
+    def _infer_and_set(module, params_holder, attr_name, run_fn):
+        """Infer conv args for a TTNN module and set them if present in parameters."""
+        if hasattr(params_holder, attr_name):
+            args = infer_ttnn_module_args(model=module, run_model=run_fn, device=None)
+            getattr(params_holder, attr_name).conv_args = args
 
-    def validate(self, output_tensor=None):
-        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])
-        )
-        output_tensor = torch.permute(output_tensor, (0, 3, 1, 2))
+    def _populate_decoder(self, torch_dec, params_dec):
+        """Warm up a single decoder (semantic or instance) to populate conv_args."""
+        if not (torch_dec and params_dec):
+            return
 
-        batch_size = output_tensor.shape[0]
+        # Synthetic tensors that match typical Panoptic-DeepLab strides
+        input_tensor = torch.randn(1, 2048, 32, 64)
+        res3_tensor = torch.randn(1, 512, 64, 128)
+        res2_tensor = torch.randn(1, 256, 128, 256)
 
-        valid_pcc = 0.97
-        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"Semantic Segmentation Head PCC check failed: {self.pcc_message}")
-        logger.info(
-            f"Panoptic DeepLab - Semantic Segmentation Head: batch_size={self.batch_size}, "
-            f"act_dtype={model_config['ACTIVATIONS_DTYPE']}, weight_dtype={model_config['WEIGHTS_DTYPE']}, "
-            f"math_fidelity={model_config['MATH_FIDELITY']}, PCC={self.pcc_message}, shape={self.output_tensor.shape}"
-        )
+        # ASPP
+        self._infer_and_set(torch_dec.aspp, params_dec, "aspp", lambda m: m(input_tensor))
+        aspp_out = torch_dec.aspp(input_tensor)
 
-        # Validate instance segmentation head outputs
-        output_tensor = self.output_tensor_2
-        output_tensor = ttnn.to_torch(output_tensor, device=self.device, mesh_composer=self.output_mesh_composer)
-        expected_shape = self.torch_output_tensor_2.shape
-        output_tensor = torch.reshape(
-            output_tensor, (expected_shape[0], expected_shape[2], expected_shape[3], expected_shape[1])
-        )
-        output_tensor = torch.permute(output_tensor, (0, 3, 1, 2))
+        # res3
+        self._infer_and_set(torch_dec.res3, params_dec, "res3", lambda m: m(aspp_out, res3_tensor))
+        res3_out = torch_dec.res3(aspp_out, res3_tensor)
 
-        batch_size = output_tensor.shape[0]
+        # res2
+        self._infer_and_set(torch_dec.res2, params_dec, "res2", lambda m: m(res3_out, res2_tensor))
+        res2_out = torch_dec.res2(res3_out, res2_tensor)
 
-        valid_pcc = 0.97
-        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"Instance Segmentation Head PCC check failed: {self.pcc_message}")
-        logger.info(
-            f"Panoptic DeepLab - Instance Segmentation Offset Head: batch_size={self.batch_size}, "
-            f"act_dtype={model_config['ACTIVATIONS_DTYPE']}, weight_dtype={model_config['WEIGHTS_DTYPE']}, "
-            f"math_fidelity={model_config['MATH_FIDELITY']}, PCC={self.pcc_message}, shape={self.output_tensor_2.shape}"
-        )
+        # heads (one or two, if present)
+        if hasattr(torch_dec, "head_1"):
+            self._infer_and_set(torch_dec.head_1, params_dec, "head_1", lambda m: m(res2_out))
+        if hasattr(torch_dec, "head_2"):
+            self._infer_and_set(torch_dec.head_2, params_dec, "head_2", lambda m: m(res2_out))
 
-        output_tensor = self.output_tensor_3
-        output_tensor = ttnn.to_torch(output_tensor, device=self.device, mesh_composer=self.output_mesh_composer)
-        expected_shape = self.torch_output_tensor_3.shape
-        output_tensor = torch.reshape(
-            output_tensor, (expected_shape[0], expected_shape[2], expected_shape[3], expected_shape[1])
-        )
-        output_tensor = torch.permute(output_tensor, (0, 3, 1, 2))
+    def _populate_all_decoders(self, torch_model, parameters):
+        if hasattr(parameters, "semantic_decoder"):
+            self._populate_decoder(torch_model.semantic_decoder, parameters.semantic_decoder)
+        if hasattr(parameters, "instance_decoder"):
+            self._populate_decoder(torch_model.instance_decoder, parameters.instance_decoder)
 
-        batch_size = output_tensor.shape[0]
+    @staticmethod
+    def _tt_to_torch_nchw(tt_tensor, device, mesh_composer, expected_shape):
+        """Convert TTNN NHWC tensor back to Torch NCHW and reshape to expected batch/shape."""
+        t = ttnn.to_torch(tt_tensor, device=device, mesh_composer=mesh_composer)
+        t = torch.reshape(t, (expected_shape[0], expected_shape[2], expected_shape[3], expected_shape[1]))
+        return torch.permute(t, (0, 3, 1, 2))
 
-        valid_pcc = 0.97
-        self.pcc_passed, self.pcc_message = check_with_pcc(self.torch_output_tensor_3, output_tensor, pcc=valid_pcc)
-        assert self.pcc_passed, logger.error(f"Instance Segmentation Head 2 PCC check failed: {self.pcc_message}")
-        logger.info(
-            f"Panoptic DeepLab - Instance Segmentation Center Head: batch_size={self.batch_size}, "
-            f"act_dtype={model_config['ACTIVATIONS_DTYPE']}, weight_dtype={model_config['WEIGHTS_DTYPE']}, "
-            f"math_fidelity={model_config['MATH_FIDELITY']}, PCC={self.pcc_message}, shape={self.output_tensor_3.shape}"
-        )
+    # --------------------------- Core runs/validation ---------------------------
 
-        return self.pcc_passed, self.pcc_message
+    def run(self):
+        self.output_tensor, self.output_tensor_2, self.output_tensor_3 = self.ttnn_model(self.input_tensor, self.device)
+        return self.output_tensor, self.output_tensor_2, self.output_tensor_3
+
+    def validate(self):
+        """Validate three heads (semantic, offsets, centers) in a uniform loop."""
+        checks = [
+            ("Semantic Segmentation Head", self.output_tensor, self.torch_output_tensor),
+            ("Instance Segmentation Offset Head", self.output_tensor_2, self.torch_output_tensor_2),
+            ("Instance Segmentation Center Head", self.output_tensor_3, self.torch_output_tensor_3),
+        ]
+
+        for name, tt_out, torch_ref in checks:
+            out = self._tt_to_torch_nchw(tt_out, self.device, self.output_mesh_composer, torch_ref.shape)
+            passed, msg = check_with_pcc(torch_ref, out, pcc=self._PCC_THRESH)
+            assert passed, logger.error(f"{name} PCC check failed: {msg}")
+
+            logger.info(
+                f"Panoptic DeepLab - {name}: batch_size={self.batch_size}, "
+                f"act_dtype={self.model_config['ACTIVATIONS_DTYPE']}, "
+                f"weight_dtype={self.model_config['WEIGHTS_DTYPE']}, "
+                f"math_fidelity={self.model_config['MATH_FIDELITY']}, "
+                f"PCC={msg}, shape={tt_out.shape}"
+            )
 
+        return True, f"All heads passed PCC ≥ {self._PCC_THRESH}"
+
+
+# --------------------------- Test config ---------------------------
 
 model_config = {
     "MATH_FIDELITY": ttnn.MathFidelity.LoFi,
@@ -260,24 +183,6 @@ def validate(self, output_tensor=None):
 
 
 @pytest.mark.parametrize("device_params", [{"l1_small_size": 16384}], indirect=True)
-@pytest.mark.parametrize(
-    "batch_size, in_channels, height, width",
-    [
-        (1, 3, 512, 1024),
-    ],
-)
-def test_panoptic_deeplab(
-    device,
-    batch_size,
-    in_channels,
-    height,
-    width,
-):
-    PanopticDeepLabTestInfra(
-        device,
-        batch_size,
-        in_channels,
-        height,
-        width,
-        model_config,
-    )
+@pytest.mark.parametrize("batch_size, in_channels, height, width", [(1, 3, 512, 1024)])
+def test_panoptic_deeplab(device, batch_size, in_channels, height, width):
+    PanopticDeepLabTestInfra(device, batch_size, in_channels, height, width, model_config)