Add half precision testing [1/n]

CHANGELOG.md

@@ -33,6 +33,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added multilabel support to `ROC` metric ([#114](https://github.com/PyTorchLightning/metrics/pull/114))
 
 
+- Added testing for `half` precision ([#77](https://github.com/PyTorchLightning/metrics/pull/77))
+
+
 ### Changed
 
 - Changed `ExplainedVariance` from storing all preds/targets to tracking 5 statistics ([#68](https://github.com/PyTorchLightning/metrics/pull/68))

docs/source/pages/overview.rst

@@ -131,6 +131,20 @@ During training and/or validation this may not be important, however it is highl
 the test dataset to only run on a single gpu or use a `join <https://pytorch.org/docs/stable/_modules/torch/nn/parallel/distributed.html#DistributedDataParallel.join>`_
 context in conjunction with DDP to prevent this behaviour. 
 
+****************************
+Metrics and 16-bit precision
+****************************
+
+Most metrics in our collection can be used with 16-bit precision (``torch.half``) tensors. However, we have found
+the following limitations:
+
+* In general ``pytorch`` had better support for 16-bit precision much earlier on GPU than CPU. Therefore, we
+  recommend that anyone that want to use metrics with half precision on CPU, upgrade to atleast pytorch v1.6
+  where support for operations such as addition, subtraction, multiplication ect. was added.
+* Some metrics does not work at all in half precision on CPU. We have explicitly stated this in their docstring,
+  but they are also listed below:
+
+  - :ref:`references/modules:PSNR` and :ref:`references/functional:psnr [func]`
 
 ******************
 Metric Arithmetics

tests/helpers/testers.py

@@ -170,6 +170,30 @@ def _functional_test(
         _assert_allclose(lightning_result, sk_result, atol=atol)
 
 
+def _assert_half_support(
+    metric_module: Metric,
+    metric_functional: Callable,
+    preds: torch.Tensor,
+    target: torch.Tensor,
+    device: str = 'cpu'
+):
+    """
+    Test if an metric can be used with half precision tensors
+
+    Args:
+        metric_module: the metric module to test
+        metric_functional: the metric functional to test
+        preds: torch tensor with predictions
+        target: torch tensor with targets
+        device: determine device, either "cpu" or "cuda"
+    """
+    p = preds[0].half().to(device) if preds[0].is_floating_point() else preds[0].to(device)
+    t = target[0].half().to(device) if target[0].is_floating_point() else target[0].to(device)
+    metric_module = metric_module.to(device)
+    assert metric_module(p, t)
+    assert metric_functional(p, t)
+
+
 class MetricTester:
     """Class used for efficiently run alot of parametrized tests in ddp mode.
     Makes sure that ddp is only setup once and that pool of processes are
@@ -288,6 +312,46 @@ def run_class_metric_test(
                 atol=self.atol,
             )
 
+    def run_precision_test_cpu(
+        self, preds: torch.Tensor, target: torch.Tensor,
+        metric_module: Metric, metric_functional: Callable, metric_args: dict = {}
+    ):
+        """ Test if an metric can be used with half precision tensors on cpu
+        Args:
+            preds: torch tensor with predictions
+            target: torch tensor with targets
+            metric_module: the metric module to test
+            metric_functional: the metric functional to test
+            metric_args: dict with additional arguments used for class initialization
+        """
+        _assert_half_support(
+            metric_module(**metric_args),
+            partial(metric_functional, **metric_args),
+            preds,
+            target,
+            device='cpu'
+        )
+
+    def run_precision_test_gpu(
+        self, preds: torch.Tensor, target: torch.Tensor,
+        metric_module: Metric, metric_functional: Callable, metric_args: dict = {}
+    ):
+        """ Test if an metric can be used with half precision tensors on gpu
+        Args:
+            preds: torch tensor with predictions
+            target: torch tensor with targets
+            metric_module: the metric module to test
+            metric_functional: the metric functional to test
+            metric_args: dict with additional arguments used for class initialization
+        """
+        _assert_half_support(
+            metric_module(**metric_args),
+            partial(metric_functional, **metric_args),
+            preds,
+            target,
+            device='cuda'
+        )
+
 
 class DummyMetric(Metric):
     name = "Dummy"

tests/regression/test_mean_error.py

@@ -24,6 +24,7 @@
 from tests.helpers.testers import BATCH_SIZE, NUM_BATCHES, MetricTester
 from torchmetrics.functional import mean_absolute_error, mean_squared_error, mean_squared_log_error
 from torchmetrics.regression import MeanAbsoluteError, MeanSquaredError, MeanSquaredLogError
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_6
 
 seed_all(42)
 
@@ -95,6 +96,20 @@ def test_mean_error_functional(self, preds, target, sk_metric, metric_class, met
             sk_metric=partial(sk_metric, sk_fn=sk_fn),
         )
 
+    @pytest.mark.skipif(
+        not _TORCH_GREATER_EQUAL_1_6,
+        reason='half support of core operations on not support before pytorch v1.6'
+    )
+    def test_mean_error_half_cpu(self, preds, target, sk_metric, metric_class, metric_functional, sk_fn):
+        if metric_class == MeanSquaredLogError:
+            # MeanSquaredLogError half + cpu does not work due to missing support in torch.log
+            pytest.xfail("MeanSquaredLogError metric does not support cpu + half precision")
+        self.run_precision_test_cpu(preds, target, metric_class, metric_functional)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason='test requires cuda')
+    def test_mean_error_half_gpu(self, preds, target, sk_metric, metric_class, metric_functional, sk_fn):
+        self.run_precision_test_gpu(preds, target, metric_class, metric_functional)
+
 
 @pytest.mark.parametrize("metric_class", [MeanSquaredError, MeanAbsoluteError, MeanSquaredLogError])
 def test_error_on_different_shape(metric_class):

tests/regression/test_psnr.py

@@ -115,6 +115,17 @@ def test_psnr_functional(self, preds, target, sk_metric, data_range, base, reduc
             metric_args=_args,
         )
 
+    # PSNR half + cpu does not work due to missing support in torch.log
+    @pytest.mark.xfail(reason="PSNR metric does not support cpu + half precision")
+    def test_psnr_half_cpu(self, preds, target, data_range, reduction, dim, base, sk_metric):
+        self.run_precision_test_cpu(preds, target, PSNR, psnr,
+                                    {"data_range": data_range, "base": base, "reduction": reduction, "dim": dim})
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason='test requires cuda')
+    def test_psnr_half_gpu(self, preds, target, data_range, reduction, dim, base, sk_metric):
+        self.run_precision_test_gpu(preds, target, PSNR, psnr,
+                                    {"data_range": data_range, "base": base, "reduction": reduction, "dim": dim})
+
 
 @pytest.mark.parametrize("reduction", ["none", "sum"])
 def test_reduction_for_dim_none(reduction):

torchmetrics/functional/regression/mean_squared_error.py

@@ -21,7 +21,8 @@
 
 def _mean_squared_error_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, int]:
     _check_same_shape(preds, target)
-    sum_squared_error = torch.sum(torch.pow(preds - target, 2))
+    diff = preds - target
+    sum_squared_error = torch.sum(diff * diff)
     n_obs = target.numel()
     return sum_squared_error, n_obs
 

torchmetrics/functional/regression/mean_squared_log_error.py

@@ -47,6 +47,10 @@ def mean_squared_log_error(preds: Tensor, target: Tensor) -> Tensor:
         >>> y = torch.tensor([0., 1, 2, 2])
         >>> mean_squared_log_error(x, y)
         tensor(0.0207)
+
+    .. note::
+        Half precision is only support on GPU for this metric
+
     """
     sum_squared_log_error, n_obs = _mean_squared_log_error_update(preds, target)
     return _mean_squared_log_error_compute(sum_squared_log_error, n_obs)

torchmetrics/functional/regression/psnr.py

@@ -41,7 +41,8 @@ def _psnr_update(
         n_obs = tensor(target.numel(), device=target.device)
         return sum_squared_error, n_obs
 
-    sum_squared_error = torch.sum(torch.pow(preds - target, 2), dim=dim)
+    diff = preds - target
+    sum_squared_error = torch.sum(diff * diff, dim=dim)
 
     if isinstance(dim, int):
         dim_list = [dim]
@@ -97,6 +98,9 @@ def psnr(
         >>> psnr(pred, target)
         tensor(2.5527)
 
+    .. note::
+        Half precision is only support on GPU for this metric
+
     """
     if dim is None and reduction != 'elementwise_mean':
         rank_zero_warn(f'The `reduction={reduction}` will not have any effect when `dim` is None.')

torchmetrics/regression/mean_squared_log_error.py

@@ -50,6 +50,9 @@ class MeanSquaredLogError(Metric):
         >>> mean_squared_log_error(preds, target)
         tensor(0.0397)
 
+    .. note::
+        Half precision is only support on GPU for this metric
+
     """
 
     def __init__(

torchmetrics/regression/psnr.py

@@ -64,6 +64,9 @@ class PSNR(Metric):
         >>> psnr(preds, target)
         tensor(2.5527)
 
+    .. note::
+        Half precision is only support on GPU for this metric
+
     """
 
     def __init__(

torchmetrics/utilities/imports.py

@@ -63,3 +63,4 @@ def _compare_version(package: str, op, version) -> bool:
 _TORCH_LOWER_1_4 = LooseVersion(torch.__version__) < LooseVersion("1.4.0")
 _TORCH_LOWER_1_5 = LooseVersion(torch.__version__) < LooseVersion("1.5.0")
 _TORCH_LOWER_1_6 = LooseVersion(torch.__version__) < LooseVersion("1.6.0")
+_TORCH_GREATER_EQUAL_1_6 = LooseVersion(torch.__version__) >= LooseVersion("1.6.0")