MaxVit model

test/test_architecture_ops.py

@@ -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__])

torchvision/models/__init__.py

@@ -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