Fixed failing tests

Author: vfdev-5

ignite/metrics/ssim.py

@@ -120,7 +120,7 @@ def reset(self) -> None:
         self._num_examples = 0
 
     def _uniform(self, kernel_size: int) -> torch.Tensor:
-        kernel = torch.zeros(kernel_size)
+        kernel = torch.zeros(kernel_size, device=self._device)
 
         start_uniform_index = max(kernel_size // 2 - 2, 0)
         end_uniform_index = min(kernel_size // 2 + 3, kernel_size)
@@ -146,10 +146,7 @@ def _gaussian_or_uniform_kernel(self, kernel_size: Sequence[int], sigma: Sequenc
 
         return torch.matmul(kernel_x.t(), kernel_y)  # (kernel_size, 1) * (1, kernel_size)
 
-    @reinit__is_reduced
-    def update(self, output: Sequence[torch.Tensor]) -> None:
-        y_pred, y = output[0].detach(), output[1].detach()
-
+    def _check_type_and_shape(self, y_pred, y):
         if y_pred.dtype != y.dtype:
             raise TypeError(
                 f"Expected y_pred and y to have the same data type. Got y_pred: {y_pred.dtype} and y: {y.dtype}."
@@ -165,6 +162,12 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
                 f"Expected y_pred and y to have BxCxHxW shape. Got y_pred: {y_pred.shape} and y: {y.shape}."
             )
 
+    @reinit__is_reduced
+    def update(self, output: Sequence[torch.Tensor]) -> None:
+        y_pred, y = output[0].detach(), output[1].detach()
+
+        self._check_type_and_shape(y_pred, y)
+
         # converts potential integer tensor to fp
         if not y.is_floating_point():
             y = y.float()
@@ -213,7 +216,15 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
         b2 = sigma_pred_sq + sigma_target_sq + self.c2
 
         ssim_idx = (a1 * a2) / (b1 * b2)
-        self._sum_of_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=self._double_dtype).sum().to(device=self._device)
+
+        # In case when ssim_idx can be MPS tensor and self._device is not MPS
+        # self._double_dtype is float64.
+        # As MPS does not support float64 we should set dtype to float32
+        double_dtype = self._double_dtype
+        if ssim_idx.device.type == "mps" and self._double_dtype == torch.float64:
+            double_dtype = torch.float32
+
+        self._sum_of_ssim += torch.mean(ssim_idx, (1, 2, 3), dtype=double_dtype).sum().to(device=self._device)
 
         self._num_examples += y.shape[0]
 

tests/ignite/conftest.py

@@ -69,19 +69,25 @@ def term_handler():
         yield  # Just pass through if SIGTERM isn't supported or we are not in the main thread
 
 
-@pytest.fixture(
-    params=[
-        "cpu",
-        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no CUDA support")),
-        pytest.param(
-            "mps", marks=pytest.mark.skipif(not is_mps_available_and_functional(), reason="Skip if no MPS support")
-        ),
-    ]
-)
+available_devices_list = [
+    "cpu",
+    pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no CUDA support")),
+    pytest.param(
+        "mps", marks=pytest.mark.skipif(not is_mps_available_and_functional(), reason="Skip if no MPS support")
+    ),
+]
+
+
+@pytest.fixture(params=available_devices_list)
 def available_device(request):
     return request.param
 
 
+@pytest.fixture(params=available_devices_list)
+def available_device2(request):
+    return request.param
+
+
 @pytest.fixture()
 def dirname():
     path = Path(tempfile.mkdtemp())

tests/ignite/metrics/gan/test_fid.py

@@ -67,8 +67,9 @@ def test_compute_fid_from_features(available_device):
     mu1, sigma1 = train_samples.mean(axis=0), cov(train_samples, rowvar=False)
     mu2, sigma2 = test_samples.mean(axis=0), cov(test_samples, rowvar=False)
 
+    tol = 1e-4 if available_device == "mps" else 1e-5
     assert (
-        pytest.approx(pytorch_fid_score.calculate_frechet_distance(mu1, sigma1, mu2, sigma2), rel=1e-5)
+        pytest.approx(pytorch_fid_score.calculate_frechet_distance(mu1, sigma1, mu2, sigma2), rel=tol)
         == fid_scorer.compute()
     )
 
@@ -142,22 +143,22 @@ def test_statistics(available_device):
     fid_scorer.update([train_samples[:5], test_samples[:5]])
     fid_scorer.update([train_samples[5:], test_samples[5:]])
 
-    mu1, sigma1 = train_samples.mean(axis=0), torch.tensor(cov(train_samples, rowvar=False))
-    mu2, sigma2 = test_samples.mean(axis=0), torch.tensor(cov(test_samples, rowvar=False))
+    mu1 = train_samples.mean(axis=0, dtype=torch.float64)
+    sigma1 = torch.tensor(cov(train_samples, rowvar=False), dtype=torch.float64)
+    mu2 = test_samples.mean(axis=0, dtype=torch.float64)
+    sigma2 = torch.tensor(cov(test_samples, rowvar=False), dtype=torch.float64)
 
     fid_mu1 = fid_scorer._train_total / fid_scorer._num_examples
     fid_sigma1 = fid_scorer._get_covariance(fid_scorer._train_sigma, fid_scorer._train_total)
 
     fid_mu2 = fid_scorer._test_total / fid_scorer._num_examples
     fid_sigma2 = fid_scorer._get_covariance(fid_scorer._test_sigma, fid_scorer._test_total)
 
-    assert torch.isclose(mu1.double(), fid_mu1.cpu()).all()
-    for cov1, cov2 in zip(sigma1, fid_sigma1):
-        assert torch.isclose(cov1.double(), cov2.cpu(), rtol=1e-04, atol=1e-04).all()
+    assert torch.allclose(mu1, fid_mu1.to(mu1))
+    assert torch.allclose(sigma1, fid_sigma1.to(sigma1), rtol=1e-04, atol=1e-04)
 
-    assert torch.isclose(mu2.double(), fid_mu2.cpu()).all()
-    for cov1, cov2 in zip(sigma2, fid_sigma2):
-        assert torch.isclose(cov1.double(), cov2.cpu(), rtol=1e-04, atol=1e-04).all()
+    assert torch.allclose(mu2, fid_mu2.to(mu2))
+    assert torch.allclose(sigma2, fid_sigma2.to(sigma2), rtol=1e-04, atol=1e-04)
 
 
 def _test_distrib_integration(device):

tests/ignite/metrics/test_ssim.py

@@ -128,18 +128,9 @@ def compare_ssim_ignite_skiimg(
     assert np.allclose(ignite_ssim, skimg_ssim, atol=precision)
 
 
-@pytest.mark.parametrize(
-    "metric_device, y_pred_device",
-    [
-        [torch.device("cpu"), torch.device("cpu")],
-        [torch.device("cpu"), torch.device("cuda")],
-        [torch.device("cuda"), torch.device("cpu")],
-        [torch.device("cuda"), torch.device("cuda")],
-    ],
-)
-def test_ssim_device(available_device, metric_device, y_pred_device):
-    if available_device == "cpu":
-        pytest.skip("This test requires a cuda device.")
+def test_ssim_device(available_device, available_device2):
+    metric_device = available_device
+    y_pred_device = available_device2
 
     data_range = 1.0
     sigma = 1.5
@@ -150,26 +141,28 @@ def test_ssim_device(available_device, metric_device, y_pred_device):
     y_pred = torch.rand(shape, device=y_pred_device)
     y = y_pred * 0.8
 
-    if metric_device == torch.device("cuda") and y_pred_device == torch.device("cpu"):
-        with pytest.warns(UserWarning):
+    if metric_device != y_pred_device and y_pred_device == "cpu":
+        with pytest.warns(
+            UserWarning,
+            match=r"y_pred tensor is on cpu device but previous computation was on another device",
+        ):
             ssim.update((y_pred, y))
     else:
         ssim.update((y_pred, y))
 
-    if metric_device == torch.device("cuda") or y_pred_device == torch.device("cuda"):
-        # A tensor will always have the device index set
-        excepted_device = torch.device("cuda:0")
+    if y_pred_device != "cpu" and metric_device == "cpu":
+        excepted_device = y_pred_device
     else:
-        excepted_device = torch.device("cpu")
+        excepted_device = metric_device
 
-    assert ssim._kernel.device == excepted_device
+    assert ssim._kernel.device.type == excepted_device
 
 
 def test_ssim_variable_batchsize(available_device):
     # Checks https://github.com/pytorch/ignite/issues/2532
     sigma = 1.5
     data_range = 1.0
-    ssim = SSIM(data_range=data_range, sigma=sigma)
+    ssim = SSIM(data_range=data_range, sigma=sigma, device=available_device)
 
     y_preds = [
         torch.rand(12, 3, 28, 28, device=available_device),
@@ -209,11 +202,14 @@ def test_ssim_variable_channel(available_device):
 @pytest.mark.parametrize(
     "dtype, precision", [(torch.bfloat16, 2e-3), (torch.float16, 4e-4), (torch.float32, 2e-5), (torch.float64, 2e-5)]
 )
-def test_cuda_ssim_dtypes(available_device, dtype, precision):
+def test_ssim_dtypes(available_device, dtype, precision):
     # Checks https://github.com/pytorch/ignite/pull/3034
     if available_device == "cpu" and dtype in [torch.float16, torch.bfloat16]:
         pytest.skip(reason=f"Unsupported dtype {dtype} on CPU device")
 
+    if available_device == "mps" and dtype in [torch.float64]:
+        pytest.skip(reason=f"Unsupported dtype {dtype} on MPS device")
+
     shape = (12, 3, 28, 28)
 
     y_pred = torch.rand(shape, device=available_device, dtype=dtype)