add autodl

AutoDL-Projects/tests/test_super_vit.py

@@ -0,0 +1,43 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
+#####################################################
+# pytest ./tests/test_super_vit.py -s               #
+#####################################################
+import unittest
+from parameterized import parameterized
+
+import torch
+from xautodl.xmodels import transformers
+from xautodl.utils.flop_benchmark import count_parameters
+
+
+class TestSuperViT(unittest.TestCase):
+    """Test the super re-arrange layer."""
+
+    def test_super_vit(self):
+        model = transformers.get_transformer("vit-base-16")
+        tensor = torch.rand((2, 3, 224, 224))
+        print("The tensor shape: {:}".format(tensor.shape))
+        # print(model)
+        outs = model(tensor)
+        print("The output tensor shape: {:}".format(outs.shape))
+
+    @parameterized.expand(
+        [
+            ["vit-cifar10-p4-d4-h4-c32", 32],
+            ["vit-base-16", 224],
+            ["vit-large-16", 224],
+            ["vit-huge-14", 224],
+        ]
+    )
+    def test_imagenet(self, name, resolution):
+        tensor = torch.rand((2, 3, resolution, resolution))
+        config = transformers.name2config[name]
+        model = transformers.get_transformer(config)
+        outs = model(tensor)
+        size = count_parameters(model, "mb", True)
+        print(
+            "{:10s} : size={:.2f}MB, out-shape: {:}".format(
+                name, size, tuple(outs.shape)
+            )
+        )