Bug fix for MS-GMSD

photosynthesis_metrics/gmsd.py

@@ -246,7 +246,7 @@ def compute_metric(self, prediction: torch.Tensor, target: torch.Tensor) -> torc
             assert prediction.size(1) == 3, "Chromatic component can be computed only for RGB images!"
 
             # Convert to YIQ color space https://en.wikipedia.org/wiki/YIQ
-            iq_weights = torch.tensor([[0.5959, -0.2746, -0.3213], [0.2115, -0.5227, 0.3112]]).t()
+            iq_weights = torch.tensor([[0.5959, -0.2746, -0.3213], [0.2115, -0.5227, 0.3112]]).t().to(prediction)
             prediction_iq = torch.matmul(prediction.permute(0, 2, 3, 1), iq_weights).permute(0, 3, 1, 2)
             target_iq = torch.matmul(target.permute(0, 2, 3, 1), iq_weights).permute(0, 3, 1, 2)
 

tests/test_gmsd.py

@@ -15,11 +15,17 @@ def target() -> torch.Tensor:
 
 
 # ================== Test class: `GMSDLoss` ==================
-def test_gmsd_loss_init() -> None:
-    try:
-        GMSDLoss()
-    except Exception as e:
-        pytest.fail(f"Unexpected error occurred: {e}")
+def test_gmsd_loss(prediction: torch.Tensor, target: torch.Tensor) -> None:
+    loss = GMSDLoss()
+    loss(prediction, target)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='No need to run test if there is no GPU.')
+def test_gmsd_loss_on_gpu(prediction: torch.Tensor, target: torch.Tensor) -> None:
+    prediction = prediction.cuda()
+    target = target.cuda()
+    loss = GMSDLoss()
+    loss(prediction, target)
 
 
 def test_gmsd_zero_for_equal_tensors(prediction: torch.Tensor):
@@ -51,18 +57,29 @@ def test_gmsd_supports_greyscale_tensors():
     loss = GMSDLoss()
     target = torch.ones(3, 1, 256, 256)
     prediction = torch.zeros(3, 1, 256, 256)
-    try:
-        loss(prediction, target)
-    except Exception as e:
-        pytest.fail(f"Unexpected error occurred: {e}")
+    loss(prediction, target)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='No need to run test if there is no GPU.')
+def test_gmsd_supports_greyscale_tensors_on_gpu():
+    loss = GMSDLoss()
+    target = torch.ones(3, 1, 256, 256).cuda()
+    prediction = torch.zeros(3, 1, 256, 256).cuda()
+    loss(prediction, target)
 
 
 # ================== Test class: `MultiScaleGMSDLoss` ==================
-def test_multi_scale_gmsd_loss_init() -> None:
-    try:
-        MultiScaleGMSDLoss()
-    except Exception as e:
-        pytest.fail(f"Unexpected error occurred: {e}")
+def test_multi_scale_gmsd_loss(prediction: torch.Tensor, target: torch.Tensor) -> None:
+    loss = MultiScaleGMSDLoss(chromatic=True)
+    loss(prediction, target)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='No need to run test if there is no GPU.')
+def test_multi_scale_gmsd_loss_on_gpu(prediction: torch.Tensor, target: torch.Tensor) -> None:
+    prediction = prediction.cuda()
+    target = target.cuda()
+    loss = MultiScaleGMSDLoss(chromatic=True)
+    loss(prediction, target)
 
 
 def test_multi_scale_gmsd_zero_for_equal_tensors(prediction: torch.Tensor):
@@ -88,10 +105,15 @@ def test_multi_scale_gmsd_supports_greyscale_tensors():
     loss = MultiScaleGMSDLoss()
     target = torch.ones(3, 1, 256, 256)
     prediction = torch.zeros(3, 1, 256, 256)
-    try:
-        loss(prediction, target)
-    except Exception as e:
-        pytest.fail(f"Unexpected error occurred: {e}")
+    loss(prediction, target)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='No need to run test if there is no GPU.')
+def test_multi_scale_gmsd_supports_greyscale_tensors_on_gpu():
+    loss = MultiScaleGMSDLoss()
+    target = torch.ones(3, 1, 256, 256).cuda()
+    prediction = torch.zeros(3, 1, 256, 256).cuda()
+    loss(prediction, target)
 
 
 def test_multi_scale_gmsd_fails_for_greyscale_tensors_chromatic_flag():
@@ -102,12 +124,17 @@ def test_multi_scale_gmsd_fails_for_greyscale_tensors_chromatic_flag():
         loss(prediction, target)
 
 
-def test_multi_scale_gmsd_supports_custom_scale_weights(prediction: torch.Tensor, target: torch.Tensor):
-    try:
-        loss = MultiScaleGMSDLoss(scale_weights=[3., 4., 2., 1., 2.])
-        loss(prediction, target)
-    except Exception as e:
-        pytest.fail(f"Unexpected error occurred: {e}")
+def test_multi_scale_gmsd_supports_custom_weights(
+        prediction: torch.Tensor, target: torch.Tensor):
+    loss = MultiScaleGMSDLoss(scale_weights=[3., 4., 2., 1., 2.])
+    loss(prediction, target)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='No need to run test if there is no GPU.')
+def test_multi_scale_gmsd_supports_custom_weights_on_gpu(
+        prediction: torch.Tensor, target: torch.Tensor):
+    loss = MultiScaleGMSDLoss(scale_weights=[3., 4., 2., 1., 2.])
+    loss(prediction, target)
 
 
 def test_multi_scale_gmsd_raise_exception_for_small_images():