clean full net test

Author: ign-febin

models/experimental/panoptic_deeplab/reference/panoptic_deeplab.py

@@ -20,13 +20,6 @@ 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()
 

models/experimental/panoptic_deeplab/tests/test_panoptic_deeplab.py

@@ -5,43 +5,52 @@
 import torch
 from loguru import logger
 import ttnn
-from ttnn.model_preprocessing import (
-    preprocess_model_parameters,
-    infer_ttnn_module_args,
-)
+from ttnn.model_preprocessing import preprocess_model_parameters
 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:
-    _seeded = False
-    _PCC_THRESH = 0.97
-
-    def __init__(self, device, batch_size, in_channels, height, width, model_config):
+    def __init__(
+        self,
+        device,
+        batch_size,
+        in_channels,
+        height,
+        width,
+        model_config,
+    ):
         super().__init__()
-        self._maybe_seed()
+        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
 
-        # Core state
+        self.pcc_passed = False
+        self.pcc_message = "call validate()?"
         self.device = device
-        self.model_config = model_config
         self.num_devices = device.get_num_devices()
-        self.batch_size, self.in_channels, self.height, self.width = (
-            batch_size,
-            in_channels,
-            height,
-            width,
-        )
+        self.batch_size = batch_size
+        self.in_channels = in_channels
+        self.height = height
+        self.width = width
         self.inputs_mesh_mapper, self.weights_mesh_mapper, self.output_mesh_composer = self.get_mesh_mappers(device)
 
-        # Torch reference model + inputs
-        torch_model = TorchPanopticDeepLab().eval()
+        # Initialize torch model
+        torch_model = TorchPanopticDeepLab()
+        torch_model = load_torch_model_state(torch_model, "panoptic_deeplab")
+
+        # Create input tensor
         input_shape = (batch_size * self.num_devices, in_channels, height, width)
         self.torch_input_tensor = torch.rand(input_shape, dtype=torch.float)
 
-        # Preprocess TTNN parameters
+        # Preprocess model parameters
         parameters = preprocess_model_parameters(
             initialize_model=lambda: torch_model,
             custom_preprocessor=create_custom_mesh_preprocessor(self.weights_mesh_mapper),
@@ -51,23 +60,27 @@ def __init__(self, device, batch_size, in_channels, height, width, model_config)
         # Populate conv_args for decoders via one small warm-up pass
         self._populate_all_decoders(torch_model, parameters)
 
-        # Run Torch once (fp32) → then bf16 for parity with TTNN
+        # Run torch model with bfloat16
         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 NHWC host tensor
+        # Convert input to TTNN format (NHWC)
         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,
         )
 
-        # TTNN model
+        # Initialize TTNN model
         logger.info("Initializing TTNN model...")
-        self.ttnn_model = TTPanopticDeepLab(parameters=parameters, model_config=model_config)
+        print("Initializing TTNN model...")
+        self.ttnn_model = TTPanopticDeepLab(
+            parameters=parameters,
+            model_config=model_config,
+        )
 
         # First run configures JIT, second run is optimized
         for phase in ("JIT configuration", "optimized"):
@@ -76,16 +89,7 @@ def __init__(self, device, batch_size, in_channels, height, width, model_config)
             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:
             return (
@@ -143,8 +147,6 @@ def _tt_to_torch_nchw(tt_tensor, device, mesh_composer, expected_shape):
         t = torch.reshape(t, (expected_shape[0], expected_shape[2], expected_shape[3], expected_shape[1]))
         return torch.permute(t, (0, 3, 1, 2))
 
-    # --------------------------- Core runs/validation ---------------------------
-
     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
@@ -157,24 +159,24 @@ def validate(self):
             ("Instance Segmentation Center Head", self.output_tensor_3, self.torch_output_tensor_3),
         ]
 
+        self._PCC_THRESH = 0.97
+
         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"act_dtype={model_config['ACTIVATIONS_DTYPE']}, "
+                f"weight_dtype={model_config['WEIGHTS_DTYPE']}, "
+                f"math_fidelity={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,
     "WEIGHTS_DTYPE": ttnn.bfloat8_b,
@@ -183,6 +185,24 @@ def validate(self):
 
 
 @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,
+    )

models/experimental/panoptic_deeplab/tt/panoptic_deeplab.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import ttnn
-from typing import Dict
 
 from models.experimental.panoptic_deeplab.tt.backbone import TTBackbone
 from models.experimental.panoptic_deeplab.tt.decoder import TTDecoder, decoder_layer_optimisations
@@ -42,9 +41,9 @@ def __init__(
 
     def __call__(
         self,
-        x: ttnn.Tensor,
+        x,
         device,
-    ) -> Dict[str, ttnn.Tensor]:
+    ):
         """
         Forward pass of TTNN Panoptic DeepLab.