Added maxvit architecture and tests

Author: TeodorPoncu

test/expect/ModelTester.test_max_vit_B_224_expect.pkl

@@ -0,0 +1,39 @@
+import unittest
+
+import pytest
+import torch
+
+from torchvision.models.maxvit import SwapAxes, WindowDepartition, WindowPartition
+
+
+class MaxvitTester(unittest.TestCase):
+    def test_maxvit_window_partition(self):
+        input_shape = (1, 3, 224, 224)
+        partition_size = 7
+
+        x = torch.randn(input_shape)
+
+        partition = WindowPartition(partition_size=7)
+        departition = WindowDepartition(partition_size=partition_size, n_partitions=(input_shape[3] // partition_size))
+
+        assert torch.allclose(x, departition(partition(x)))
+
+    def test_maxvit_grid_partition(self):
+        input_shape = (1, 3, 224, 224)
+        partition_size = 7
+
+        x = torch.randn(input_shape)
+        partition = torch.nn.Sequential(
+            WindowPartition(partition_size=(input_shape[3] // partition_size)),
+            SwapAxes(-2, -3),
+        )
+        departition = torch.nn.Sequential(
+            SwapAxes(-2, -3),
+            WindowDepartition(partition_size=(input_shape[3] // partition_size), n_partitions=partition_size),
+        )
+
+        assert torch.allclose(x, departition(partition(x)))
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

test/expect/ModelTester.test_max_vit_L_224_expect.pkl

@@ -594,6 +594,14 @@ def test_vitc_models(model_fn, dev):
     test_classification_model(model_fn, dev)
 
 
+@pytest.mark.parametrize(
+    "model_fn", [models.max_vit_T_224, models.max_vit_S_224, models.max_vit_B_224, models.max_vit_L_224]
+)
+@pytest.mark.parametrize("dev", cpu_and_gpu())
+def test_max_vit(model_fn, dev):
+    test_classification_model(model_fn, dev)
+
+
 @pytest.mark.parametrize("model_fn", list_model_fns(models))
 @pytest.mark.parametrize("dev", cpu_and_gpu())
 def test_classification_model(model_fn, dev):

test/expect/ModelTester.test_max_vit_S_224_expect.pkl

@@ -13,5 +13,6 @@
 from .vgg import *
 from .vision_transformer import *
 from .swin_transformer import *
+from .maxvit import *
 from . import detection, optical_flow, quantization, segmentation, video
 from ._api import get_model, get_model_weights, get_weight, list_models