PSNR implementation

pytorch_lightning/metrics/__init__.py

@@ -5,6 +5,7 @@
     RMSE,
     MAE,
     RMSLE,
+    PSNR
 )
 from pytorch_lightning.metrics.classification import (
     Accuracy,
@@ -50,6 +51,7 @@
     'MSE',
     'RMSE',
     'MAE',
-    'RMSLE'
+    'RMSLE',
+    'PSNR'
 ]
 __all__ = __regression_metrics + __classification_metrics + ['SklearnMetric']

pytorch_lightning/metrics/functional/regression.py

@@ -0,0 +1,38 @@
+import torch
+from torch.nn import functional as F
+
+from pytorch_lightning.metrics.functional.reduction import reduce
+
+def psnr(
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    data_range: float = None,
+    base: float = 10.0,
+    reduction: str = 'elementwise_mean'
+) -> torch.Tensor:
+    """
+    Computes the peak signal-to-noise ratio metric
+
+    Args:
+        pred: estimated signal
+        target: groun truth signal
+        data_range: the range of the data. If None, it is determined from the data (max - min).
+        base: a base of a logarithm to use (default: 10)
+        reduction: method for reducing psnr (default: takes the mean)
+
+    Example:
+
+        >>> pred = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
+        >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
+        >>> metric = PSNR()
+        >>> metric(pred, target)
+        tensor([2.5527])
+        """
+
+    if data_range is None:
+        data_range = max(target.max() - target.min(), pred.max() - pred.min())
+    else:
+        data_range = torch.tensor(data_range)
+    mse = F.mse_loss(pred.view(-1), target.view(-1), reduction=reduction)
+    psnr_base_e = 2 * torch.log(data_range) - torch.log(mse)
+    return psnr_base_e * (10 / torch.log(base))

pytorch_lightning/metrics/regression.py

@@ -2,7 +2,7 @@
 import torch
 from pytorch_lightning.metrics.metric import Metric
 
-__all__ = ['MSE', 'RMSE', 'MAE', 'RMSLE']
+__all__ = ['MSE', 'RMSE', 'MAE', 'RMSLE', 'PSNR']
 
 
 class MSE(Metric):
@@ -187,3 +187,38 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             A Tensor with the rmsle loss.
         """
         return F.mse_loss(torch.log(pred + 1), torch.log(target + 1), self.reduction)
+
+
+class PSNR(Metric):
+    """
+    Computes the peak signal-to-noise ratio metric
+    """
+
+    def __init__(self, data_range: float = None, base: int = 10):
+        """
+        Args:
+            data_range: the range of the data. If None, it is determined from the data (max - min).
+            base: a base of a logarithm to use (default: 10)
+
+
+        Example:
+
+            >>> pred = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
+            >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
+            >>> metric = PSNR()
+            >>> metric(pred, target)
+            tensor([2.5527])
+        """
+        super().__init__(name='psnr')
+        self.data_range = data_range
+        self.base = torch.tensor([float(base)])
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        if self.data_range is None:
+            data_range = max(target.max() - target.min(), pred.max() - pred.min())
+        else:
+            data_range = torch.tensor(float(self.data_range))
+        mse = F.mse_loss(pred.view(-1), target.view(-1))
+        # numerical precision tricks
+        psnr_base_e = 2 * torch.log(data_range) - torch.log(mse)
+        return psnr_base_e * (10 / torch.log(self.base))  # change the logarithm basis

requirements/test.txt

@@ -7,7 +7,7 @@ flake8
 flake8-black
 check-manifest
 twine==1.13.0
-
+scikit-image
 black==19.10b0
 pre-commit>=1.0
 

tests/metrics/test_regression.py

@@ -1,8 +1,10 @@
 import pytest
 import torch
+from skimage.metrics import peak_signal_noise_ratio as ski_psnr
+import numpy as np
 
 from pytorch_lightning.metrics.regression import (
-    MAE, MSE, RMSE, RMSLE
+    MAE, MSE, RMSE, RMSLE, PSNR
 )
 
 
@@ -64,3 +66,44 @@ def test_rmsle(pred, target, exp):
 
     assert isinstance(score, torch.Tensor)
     assert pytest.approx(score.item(), rel=1e-3) == exp
+
+@pytest.mark.parametrize(['pred', 'target', 'exp'], [
+    pytest.param(
+        [0., 1., 2., 3.],
+        [0., 1., 2., 2.],
+        ski_psnr(np.array([0., 1., 2., 3.]), np.array([0., 1., 2., 2.]), data_range=3)
+    ),
+    pytest.param(
+        [4., 3., 2., 1.,],
+        [1., 4., 3., 2.,],
+        ski_psnr(np.array([4., 3., 2., 1.]), np.array([1., 4., 3., 2.]), data_range=3)
+    )
+])
+def test_psnr(pred, target, exp):
+    psnr = PSNR()
+    assert psnr.name == 'psnr'
+    score = psnr(pred=torch.tensor(pred),
+                 target=torch.tensor(target))
+    assert isinstance(score, torch.Tensor)
+    assert pytest.approx(score.item(), rel=1e-3) == exp
+
+@pytest.mark.parametrize(['pred', 'target', 'exp'], [
+    pytest.param(
+        [0., 1., 2., 3.],
+        [0., 1., 2., 2.],
+        ski_psnr(np.array([0., 1., 2., 3.]), np.array([0., 1., 2., 2.]), data_range=4) * np.log(10)
+    ),
+    pytest.param(
+        [4., 3., 2., 1.,],
+        [1., 4., 3., 2.,],
+        ski_psnr(np.array([4., 3., 2., 1.]), np.array([1., 4., 3., 2.]), data_range=4) * np.log(10)
+    )
+])
+
+def test_psnr_base_e_wider_range(pred, target, exp):
+    psnr = PSNR(data_range=4, base=2.718281828459045)
+    assert psnr.name == 'psnr'
+    score = psnr(pred=torch.tensor(pred),
+                 target=torch.tensor(target))
+    assert isinstance(score, torch.Tensor)
+    assert pytest.approx(score.item(), rel=1e-3) == exp