tests update

Author: ign-krishnanand

models/experimental/SSR/tests/test_RHAG.py

@@ -93,12 +93,7 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
 @pytest.mark.parametrize(
     "batch_size, height, width, dim, num_heads, window_size, depth, overlap_ratio, mlp_ratio, resi_connection",
     [
-        # (1, 32, 32, 180, 6, 16, 2, 0.5, 2.0, "1conv"),      # Standard configuration with conv
-        (1, 64, 64, 180, 6, 16, 6, 0.5, 2.0, "1conv"),  # Standard configuration with conv
-        # (1, 64, 64, 180, 6, 16, 2, 0.5, 2.0, "1conv"),      # Standard configuration with conv
-        # (1, 32, 32, 96, 3, 8, 3, 0.25, 4.0, "identity"),  # Identity connection
-        # (2, 64, 64, 180, 6, 16, 1, 0.5, 2.0, "1conv"),      # Batch size 2
-        # (1, 128, 128, 192, 6, 16, 2, 0.75, 3.0, "1conv"),   # Larger resolution
+        (1, 64, 64, 180, 6, 16, 6, 0.5, 2, "1conv"),  # SSR config
     ],
 )
 @pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)

models/experimental/SSR/tests/test_atten_blocks.py

@@ -37,11 +37,7 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
 @pytest.mark.parametrize(
     "batch_size, height, width, dim, num_heads, window_size, depth, overlap_ratio, mlp_ratio",
     [
-        (1, 64, 64, 180, 6, 16, 2, 0.5, 2.0),  # Standard configuration
-        # (1, 32, 32, 96, 3, 8, 3, 0.25, 4.0),  # Smaller resolution, more blocks
-        # (2, 64, 64, 180, 6, 16, 1, 0.5, 2.0),  # Batch size 2, single block
-        # (1, 128, 128, 192, 6, 16, 2, 0.75, 3.0),  # Larger resolution
-        # (2, 64, 64, 180, 6, 16, 6, 0.5, 2),  # Network config
+        (1, 64, 64, 180, 6, 16, 6, 0.5, 2),  # SSR config
     ],
 )
 @pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)

models/experimental/SSR/tests/test_patch_embed_tile_selection.py renamed to models/experimental/SSR/tests/test_patch_embed_tile_refinement.py

@@ -131,11 +131,8 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
 @pytest.mark.parametrize(
     "batch_size, img_size, patch_size, in_chans, embed_dim, norm_layer",
     [
-        (1, 224, 16, 3, 768, None),  # Standard ViT-Base config
-        (2, 256, 4, 3, 96, None),  # Smaller embedding
-        (1, 32, 4, 180, 180, None),  # Your SSR config
-        # (1, 64, 8, 3, 192, nn.LayerNorm),     # With normalization
-        (4, 128, 16, 3, 384, None),  # Batch size 4
+        (1, 64, 2, 3, 180, None),  # TR blk test
+        (1, 64, 4, 3, 180, None),  # HAT blk test
     ],
 )
 @pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
@@ -200,11 +197,13 @@ def test_patch_embed_simple(device, batch_size, img_size, patch_size, in_chans,
     logger.info(pcc_message)
 
     if does_pass:
-        logger.info("Simple PatchEmbed Passed!")
+        logger.info("TR PatchEmbed Passed!")
     else:
-        logger.warning("Simple PatchEmbed Failed!")
+        logger.warning("TR PatchEmbed Failed!")
 
     assert does_pass, f"PCC check failed: {pcc_message}"
     assert (
         ref_output.shape == tt_torch_output.shape
     ), f"Shape mismatch: ref {ref_output.shape} vs ttnn {tt_torch_output.shape}"
+
+    ttnn.close_device(device)

models/experimental/SSR/tests/test_patch_unembed.py

@@ -1,18 +1,22 @@
 import pytest
 import torch
 import ttnn
-from tests.ttnn.utils_for_testing import assert_with_pcc
+from loguru import logger
+from models.utility_functions import comp_pcc
+
 from models.utility_functions import torch_random
 
 from models.experimental.SSR.reference.SSR.model.tile_refinement import PatchUnEmbed
 from models.experimental.SSR.tt.patch_unembed import TTPatchUnEmbed
 
 
-@pytest.mark.parametrize("batch_size", [1, 2, 8])
-@pytest.mark.parametrize("img_size", [64])
-@pytest.mark.parametrize("patch_size", [1])
-@pytest.mark.parametrize("in_chans", [180])
-@pytest.mark.parametrize("embed_dim", [180])
+@pytest.mark.parametrize(
+    "batch_size, img_size, patch_size, in_chans, embed_dim",
+    [
+        # (1, 64, 4, 3, 180),   # TR blk test
+        (1, 64, 2, 3, 180),  # HAT blk test
+    ],
+)
 def test_tt_patch_unembed(device, batch_size, img_size, patch_size, in_chans, embed_dim):
     """Test TTPatchUnEmbed against PyTorch reference implementation"""
     torch.manual_seed(0)
@@ -37,17 +41,29 @@ def test_tt_patch_unembed(device, batch_size, img_size, patch_size, in_chans, em
 
     # Convert input to TTNN format
     ttnn_input = ttnn.from_torch(
-        torch_input, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.DRAM_MEMORY_CONFIG
+        torch_input, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.L1_MEMORY_CONFIG
     )
 
     # Run TTNN model
     ttnn_output = tt_model(ttnn_input, patches_resolution)
     ttnn_output_torch = ttnn.to_torch(ttnn_output)
 
+    # Compare outputs
+    does_pass, pcc_message = comp_pcc(torch_output, ttnn_output_torch, 0.99)
+
+    logger.info(f"Reference output shape: {torch_output.shape}")
+    logger.info(f"TTNN output shape: {ttnn_output_torch.shape}")
+    logger.info(pcc_message)
+
+    if does_pass:
+        logger.info("TR PatchEmbed Passed!")
+    else:
+        logger.warning("TR PatchEmbed Failed!")
+
+    assert does_pass, f"PCC check failed: {pcc_message}"
     # Assert shapes match
     assert (
         torch_output.shape == ttnn_output_torch.shape
     ), f"Shape mismatch: {torch_output.shape} vs {ttnn_output_torch.shape}"
 
-    # Assert values match with PCC
-    assert_with_pcc(torch_output, ttnn_output_torch, 0.999)
+    ttnn.close_device(device)

models/experimental/SSR/tests/test_tile_refinement.py

@@ -118,7 +118,7 @@ def custom_preprocessor(torch_model, name, ttnn_module_args):
         # Test configuration - adjust based on your requirements
         # (64, 1, 96, (2, 2), (6, 6), 7, 4, 2, (1, 3, 64, 64)),
         # (64, 1, 180, (2, 2, 2), (6, 6, 6), 16, 2, 4, (1, 3, 64, 64)),
-        (64, 1, 180, (6, 6, 6, 6, 6, 6), (6, 6, 6, 6, 6, 6), 16, 2, 4, (3, 3, 64, 64)),
+        (64, 2, 180, (6, 6, 6, 6, 6, 6), (6, 6, 6, 6, 6, 6), 16, 2, 4, (3, 3, 64, 64)),
     ],
 )
 def test_tile_refinement(

models/experimental/SSR/tt/atten_blocks.py

@@ -73,11 +73,8 @@ def __init__(
             mlp_ratio=mlp_ratio,
         )
 
-        # Downsample layer (if provided)
         self.downsample = None
         if downsample is not None:
-            # Note: You'll need to implement the downsample layer in TTNN
-            # This depends on what type of downsampling is used
             self.downsample = downsample
 
     def forward(self, x, x_size, params):

models/experimental/SSR/tt/patch_embed_tile_refinement.py

@@ -36,9 +36,6 @@ def __init__(
         self.patches_resolution = [self.img_size[0] // self.patch_size[0], self.img_size[1] // self.patch_size[1]]
         self.num_patches = self.patches_resolution[0] * self.patches_resolution[1]
 
-        self.in_chans = in_chans
-        self.embed_dim = embed_dim
-        self.device = device
         self.memory_config = memory_config or ttnn.DRAM_MEMORY_CONFIG
 
         # Store normalization parameters if provided

models/experimental/SSR/tt/patch_unembed.py

@@ -8,16 +8,6 @@ class TTPatchUnEmbed(LightweightModule):
     def __init__(self, mesh_device, img_size=224, patch_size=4, in_chans=3, embed_dim=96):
         super().__init__()
 
-        self.mesh_device = mesh_device
-
-        # Convert to tuples like in the original
-        self.img_size = (img_size, img_size)
-        self.patch_size = (patch_size, patch_size)
-
-        self.patches_resolution = [self.img_size[0] // self.patch_size[0], self.img_size[1] // self.patch_size[1]]
-        self.num_patches = self.patches_resolution[0] * self.patches_resolution[1]
-
-        self.in_chans = in_chans
         self.embed_dim = embed_dim
 
     def forward(self, x, x_size):
@@ -27,6 +17,6 @@ def forward(self, x, x_size):
         x = ttnn.permute(x, (0, 2, 1))  # (batch_size, embed_dim, num_patches)
 
         # Reshape to spatial dimensions
-        x = ttnn.reshape(x, (batch_size, self.embed_dim, x_size[0], x_size[1]))
+        x = ttnn.reshape(x, (batch_size, self.embed_dim, x_size[0], x_size[1]), memory_config=ttnn.L1_MEMORY_CONFIG)
 
         return x