Add compile test for models

Author: qubvel

tests/models/base.py

@@ -14,13 +14,15 @@
     default_device,
     slow_test,
     requires_torch_greater_or_equal,
+    check_run_test_on_diff_or_main,
 )
 
 
 class BaseModelTester(unittest.TestCase):
     test_encoder_name = (
         "tu-test_resnet.r160_in1k" if has_timm_test_models else "resnet18"
     )
+    files_for_diff = [r".*"]
 
     # should be overriden
     test_model_type = None
@@ -234,3 +236,21 @@ def test_preserve_forward_output(self):
         is_close = torch.allclose(output, output_tensor, atol=5e-2)
         max_diff = torch.max(torch.abs(output - output_tensor))
         self.assertTrue(is_close, f"Max diff: {max_diff}")
+
+    @pytest.mark.compile
+    def test_compile(self):
+        if not check_run_test_on_diff_or_main(self.files_for_diff):
+            self.skipTest("No diff and not on `main`.")
+
+        sample = self._get_sample(
+            batch_size=self.default_batch_size,
+            num_channels=self.default_num_channels,
+            height=self.default_height,
+            width=self.default_width,
+        ).to(default_device)
+
+        model = self.get_default_model()
+        compiled_model = torch.compile(model, fullgraph=True, dynamic=True)
+
+        with torch.inference_mode():
+            compiled_model(sample)

tests/models/test_deeplab.py

@@ -3,11 +3,13 @@
 
 class TestDeeplabV3Model(base.BaseModelTester):
     test_model_type = "deeplabv3"
+    files_for_diff = [r"decoders/deeplabv3/", r"base/"]
 
     default_batch_size = 2
 
 
 class TestDeeplabV3PlusModel(base.BaseModelTester):
     test_model_type = "deeplabv3plus"
+    files_for_diff = [r"decoders/deeplabv3plus/", r"base/"]
 
     default_batch_size = 2

tests/models/test_fpn.py

@@ -3,3 +3,4 @@
 
 class TestFpnModel(base.BaseModelTester):
     test_model_type = "fpn"
+    files_for_diff = [r"decoders/fpn/", r"base/"]

tests/models/test_linknet.py

@@ -3,3 +3,4 @@
 
 class TestLinknetModel(base.BaseModelTester):
     test_model_type = "linknet"
+    files_for_diff = [r"decoders/linknet/", r"base/"]

tests/models/test_manet.py

@@ -3,3 +3,4 @@
 
 class TestManetModel(base.BaseModelTester):
     test_model_type = "manet"
+    files_for_diff = [r"decoders/manet/", r"base/"]

tests/models/test_pan.py

@@ -3,6 +3,7 @@
 
 class TestPanModel(base.BaseModelTester):
     test_model_type = "pan"
+    files_for_diff = [r"decoders/pan/", r"base/"]
 
     default_batch_size = 2
     default_height = 128

tests/models/test_psp.py

@@ -3,5 +3,6 @@
 
 class TestPspModel(base.BaseModelTester):
     test_model_type = "pspnet"
+    files_for_diff = [r"decoders/pspnet/", r"base/"]
 
     default_batch_size = 2

tests/models/test_segformer.py

@@ -8,6 +8,7 @@
 
 class TestSegformerModel(base.BaseModelTester):
     test_model_type = "segformer"
+    files_for_diff = [r"decoders/segformer/", r"base/"]
 
     @slow_test
     @requires_torch_greater_or_equal("2.0.1")

tests/models/test_unet.py

@@ -3,3 +3,4 @@
 
 class TestUnetModel(base.BaseModelTester):
     test_model_type = "unet"
+    files_for_diff = [r"decoders/unet/", r"base/"]

tests/models/test_unetplusplus.py

@@ -3,3 +3,4 @@
 
 class TestUnetPlusPlusModel(base.BaseModelTester):
     test_model_type = "unetplusplus"
+    files_for_diff = [r"decoders/unetplusplus/", r"base/"]