Add device

Author: qubvel

Makefile

@@ -10,7 +10,7 @@ test: .venv
 	.venv/bin/pytest -v -rsx -n 2 tests/ -k "not logits_match"
 
 test_all: .venv
-	.venv/bin/pytest -v -rsx -n 2 tests/
+	RUN_SLOW=1 .venv/bin/pytest -v -rsx -n 2 tests/
 
 table:
 	.venv/bin/python misc/generate_table.py

pyproject.toml

@@ -40,6 +40,8 @@ docs = [
 ]
 test = [
     'pytest',
+    'pytest-cov',
+    'pytest-xdist',
     'ruff',
 ]
 

tests/models/base.py

@@ -9,7 +9,12 @@
 import torch
 import segmentation_models_pytorch as smp
 
-from tests.utils import has_timm_test_models, slow_test, requires_torch_greater_or_equal
+from tests.utils import (
+    has_timm_test_models,
+    default_device,
+    slow_test,
+    requires_torch_greater_or_equal,
+)
 
 
 class BaseModelTester(unittest.TestCase):
@@ -58,10 +63,10 @@ def test_forward_backward(self):
             num_channels=self.default_num_channels,
             height=self.default_height,
             width=self.default_width,
-        )
+        ).to(default_device)
         model = smp.create_model(
             arch=self.model_type, encoder_name=self.test_encoder_name
-        )
+        ).to(default_device)
 
         # check default in_channels=3
         output = model(sample)
@@ -93,13 +98,13 @@ def test_in_channels_and_depth_and_out_classes(
             in_channels=in_channels,
             classes=classes,
             **kwargs,
-        )
+        ).to(default_device)
         sample = self._get_sample(
             batch_size=self.default_batch_size,
             num_channels=in_channels,
             height=self.default_height,
             width=self.default_width,
-        )
+        ).to(default_device)
 
         # check in channels correctly set
         with torch.no_grad():
@@ -117,7 +122,7 @@ def test_classification_head(self):
                 "dropout": 0.5,
                 "activation": "sigmoid",
             },
-        )
+        ).to(default_device)
 
         self.assertIsNotNone(model.classification_head)
         self.assertIsInstance(model.classification_head[0], torch.nn.AdaptiveAvgPool2d)
@@ -132,7 +137,7 @@ def test_classification_head(self):
             num_channels=self.default_num_channels,
             height=self.default_height,
             width=self.default_width,
-        )
+        ).to(default_device)
 
         with torch.no_grad():
             _, cls_probs = model(sample)
@@ -144,14 +149,14 @@ def test_save_load_with_hub_mixin(self):
         # instantiate model
         model = smp.create_model(
             arch=self.model_type, encoder_name=self.test_encoder_name
-        )
+        ).to(default_device)
 
         # save model
         with tempfile.TemporaryDirectory() as tmpdir:
             model.save_pretrained(
                 tmpdir, dataset="test_dataset", metrics={"my_awesome_metric": 0.99}
             )
-            restored_model = smp.from_pretrained(tmpdir)
+            restored_model = smp.from_pretrained(tmpdir).to(default_device)
             with open(os.path.join(tmpdir, "README.md"), "r") as f:
                 readme = f.read()
 
@@ -161,7 +166,7 @@ def test_save_load_with_hub_mixin(self):
             num_channels=self.default_num_channels,
             height=self.default_height,
             width=self.default_width,
-        )
+        ).to(default_device)
 
         with torch.no_grad():
             output = model(sample)
@@ -178,7 +183,7 @@ def test_save_load_with_hub_mixin(self):
     @requires_torch_greater_or_equal("2.0.1")
     @pytest.mark.logits_match
     def test_preserve_forward_output(self):
-        model = smp.from_pretrained(self.hub_checkpoint).eval()
+        model = smp.from_pretrained(self.hub_checkpoint).eval().to(default_device)
 
         input_tensor_path = hf_hub_download(
             repo_id=self.hub_checkpoint, filename="input-tensor.pth"
@@ -188,12 +193,14 @@ def test_preserve_forward_output(self):
         )
 
         input_tensor = torch.load(input_tensor_path, weights_only=True)
+        input_tensor = input_tensor.to(default_device)
         output_tensor = torch.load(output_tensor_path, weights_only=True)
+        output_tensor = output_tensor.to(default_device)
 
         with torch.no_grad():
             output = model(input_tensor)
 
         self.assertEqual(output.shape, output_tensor.shape)
-        is_close = torch.allclose(output, output_tensor, atol=1e-3)
+        is_close = torch.allclose(output, output_tensor, atol=1e-2)
         max_diff = torch.max(torch.abs(output - output_tensor))
         self.assertTrue(is_close, f"Max diff: {max_diff}")