[FX][Tests] Use smaller swin transformer to speed up pre-commit

tests/torch/fx/test_models.py

@@ -45,6 +45,7 @@
 from tests.torch.test_models.synthetic import MultiBranchesConnectedModel
 from tests.torch.test_models.synthetic import ShortTransformer
 from tests.torch.test_models.synthetic import SplitCatModel
+from tests.torch.test_models.synthetic import SwinV2SingleBlock
 from tests.torch.test_models.synthetic import TopKModel
 from tests.torch.test_models.synthetic import YOLO11N_SDPABlock
 
@@ -73,7 +74,7 @@ def torchvision_model_case(model_id: str, input_shape: tuple[int,]):
     torchvision_model_case("resnet18", (1, 3, 224, 224)),
     torchvision_model_case("mobilenet_v3_small", (1, 3, 224, 224)),
     torchvision_model_case("vit_b_16", (1, 3, 224, 224)),
-    torchvision_model_case("swin_v2_t", (1, 3, 224, 224)),
+    ModelCase(SwinV2SingleBlock, "swin_v2_single_block", SwinV2SingleBlock.INPUT_SHAPE),
     ModelCase(test_models.UNet, "unet", [1, 3, 224, 224]),
     ModelCase(partial(ShortTransformer, 5, 10), "synthetic_transformer", [5]),
     ModelCase(YOLO11N_SDPABlock, "yolo11n_sdpa_block", YOLO11N_SDPABlock.INPUT_SIZE),
@@ -173,11 +174,11 @@ def test_model(test_case: ModelCase, regen_ref_data: bool):
         [Dim.AUTO, Dim.STATIC, Dim.STATIC, Dim.STATIC],  # This ViT Model is not eligible for dynamic shape capability
     ),
     (
-        torchvision_model_case("swin_v2_t", (1, 3, 224, 224)),
+        ModelCase(SwinV2SingleBlock, "swin_v2_single_block", SwinV2SingleBlock.INPUT_SHAPE),
         {"model_type": nncf.ModelType.TRANSFORMER},
         [
-            (130, 130),
-            (77, 77),
+            (50, 50),
+            (29, 29),
         ],
         [Dim.AUTO, Dim.STATIC, Dim.AUTO, Dim.AUTO],
     ),

tests/torch/test_models/synthetic.py

@@ -16,6 +16,9 @@
 from torch.nn import BatchNorm2d
 from torch.nn import Dropout
 from torch.nn import Parameter
+from torchvision.models.swin_transformer import PatchMergingV2
+from torchvision.models.swin_transformer import SwinTransformer
+from torchvision.models.swin_transformer import SwinTransformerBlockV2
 from torchvision.transforms.functional import normalize
 
 from tests.torch.helpers import create_bn
@@ -690,3 +693,26 @@ def __init__(self) -> None:
     def forward(self, x):
         # input_shape = [1, 3, 32, 32]
         return self.depthwise_conv(x)
+
+
+class SwinV2SingleBlock(nn.Module):
+    """Minimal Swin Transformer V2 model with a single block for testing."""
+
+    INPUT_SHAPE = (1, 3, 224, 224)
+
+    def __init__(self):
+        super().__init__()
+        self.model = SwinTransformer(
+            patch_size=[4, 4],
+            embed_dim=96,
+            depths=[1, 1, 1, 1],
+            num_heads=[3, 6, 12, 24],
+            window_size=[8, 8],
+            stochastic_depth_prob=0.2,
+            num_classes=10,
+            block=SwinTransformerBlockV2,
+            downsample_layer=PatchMergingV2,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)