Refactor test do not skip

Author: qubvel

segmentation_models_pytorch/base/model.py

@@ -11,6 +11,10 @@
 class SegmentationModel(torch.nn.Module, SMPHubMixin):
     """Base class for all segmentation models."""
 
+    _is_torch_scriptable = True
+    _is_torch_exportable = True
+    _is_torch_compilable = True
+
     # if model supports shape not divisible by 2 ^ n set to False
     requires_divisible_input_shape = True
 

segmentation_models_pytorch/decoders/unetplusplus/model.py

@@ -56,6 +56,8 @@ class UnetPlusPlus(SegmentationModel):
 
     """
 
+    _is_torch_scriptable = False
+
     @supports_config_loading
     def __init__(
         self,

segmentation_models_pytorch/decoders/upernet/model.py

@@ -48,6 +48,8 @@ class UPerNet(SegmentationModel):
 
     """
 
+    _is_torch_scriptable = False
+
     @supports_config_loading
     def __init__(
         self,

segmentation_models_pytorch/encoders/efficientnet.py

@@ -35,6 +35,9 @@
 class EfficientNetEncoder(EfficientNet, EncoderMixin):
     _is_torch_scriptable = False
 
+    # works with torch 2.4.0, but not with torch 2.5.1
+    _is_torch_compilable = False
+
     def __init__(
         self,
         stage_idxs: List[int],

tests/encoders/base.py

@@ -33,10 +33,6 @@ class BaseEncoderTester(unittest.TestCase):
     depth_to_test = [3, 4, 5]
     strides_to_test = [8, 16]  # 32 is a default one
 
-    # enable/disable tests
-    do_test_torch_compile = True
-    do_test_torch_export = True
-
     def get_tiny_encoder(self):
         return smp.encoders.get_encoder(self.encoder_names[0], encoder_weights=None)
 
@@ -208,28 +204,25 @@ def test_dilated(self):
 
     @pytest.mark.compile
     def test_compile(self):
-        if not self.do_test_torch_compile:
-            self.skipTest(
-                f"torch_compile test is disabled for {self.encoder_names[0]}."
-            )
-
         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().to(default_device)
 
-        encoder = self.get_tiny_encoder().eval().to(default_device)
+        encoder = self.get_tiny_encoder()
+        encoder = encoder.eval().to(default_device)
+
         compiled_encoder = torch.compile(encoder, fullgraph=True, dynamic=True)
 
-        with torch.inference_mode():
+        if encoder._is_torch_compilable:
             compiled_encoder(sample)
+        else:
+            with self.assertRaises(Exception):
+                compiled_encoder(sample)
 
     @pytest.mark.torch_export
     @requires_torch_greater_or_equal("2.4.0")
     def test_torch_export(self):
-        if not self.do_test_torch_export:
-            self.skipTest(f"torch_export test is disabled for {self.encoder_names[0]}.")
-
         if not check_run_test_on_diff_or_main(self.files_for_diff):
             self.skipTest("No diff and not on `main`.")
 

tests/encoders/test_pretrainedmodels_encoders.py

@@ -12,10 +12,6 @@ class TestDPNEncoder(base.BaseEncoderTester):
     )
     files_for_diff = ["encoders/dpn.py"]
 
-    # works with torch 2.4.0, but not with torch 2.5.1
-    # dynamo error, probably on Sequential + OrderedDict
-    do_test_torch_export = False
-
     def get_tiny_encoder(self):
         params = {
             "stage_idxs": [2, 3, 4, 6],
@@ -45,10 +41,6 @@ class TestInceptionV4Encoder(base.BaseEncoderTester):
     files_for_diff = ["encoders/inceptionv4.py"]
     supports_dilated = False
 
-    # works with torch 2.4.0, but not with torch 2.5.1
-    # dynamo error, probably on Sequential + OrderedDict
-    do_test_torch_export = False
-
 
 class TestSeNetEncoder(base.BaseEncoderTester):
     encoder_names = (

tests/encoders/test_smp_encoders.py

@@ -62,6 +62,3 @@ class TestEfficientNetEncoder(base.BaseEncoderTester):
         ]
     )
     files_for_diff = ["encoders/efficientnet.py"]
-
-    # torch_compile is not supported for efficientnet encoders
-    do_test_torch_compile = False

tests/encoders/test_timm_ported_encoders.py

@@ -26,9 +26,6 @@ class TestTimmEfficientNetEncoder(base.BaseEncoderTester):
     )
     files_for_diff = ["encoders/timm_efficientnet.py"]
 
-    # works with torch 2.4.0, but not with torch 2.5.1
-    do_test_torch_export = False
-
 
 class TestTimmGERNetEncoder(base.BaseEncoderTester):
     encoder_names = (

tests/models/base.py

@@ -263,6 +263,11 @@ def test_torch_script(self):
         model = self.get_default_model()
         model.eval()
 
+        if not model._is_torch_scriptable:
+            with self.assertRaises(RuntimeError):
+                scripted_model = torch.jit.script(model)
+            return
+
         scripted_model = torch.jit.script(model)
 
         with torch.inference_mode():