Add functional regression metrics (#2492)

* Add functional regression metrics * add functional tests * add docs * changelog * init * pep8 * docs * docs * setup docs * docs * Apply suggestions from code review * Apply suggestions from code review * typo Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: Jirka <jirka@pytorchlightning.ai>
Lightning-AI · Jul 9, 2020 · 6f4a488 · 6f4a488
1 parent 4bbcfa0
commit 6f4a488
Show file tree

Hide file tree

Showing 11 changed files with 380 additions and 184 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added a PSNR metric: peak signal-to-noise ratio ([#2483](https://github.com/PyTorchLightning/pytorch-lightning/pull/2483))
 
+- Added functional regression metrics ([#2492](https://github.com/PyTorchLightning/pytorch-lightning/pull/2492))
+
 ### Changed
 
 

diff --git a/docs/source/introduction_guide.rst b/docs/source/introduction_guide.rst
@@ -22,7 +22,7 @@ quickly becomes onerous and starts distracting from the core research code.
 Goal of this guide
 ------------------
 This guide walks through the major parts of the library to help you understand
-what each parts does. But at the end of the day, you write the same PyTorch code... just organize it
+what each part does. But at the end of the day, you write the same PyTorch code... just organize it
 into the LightningModule template which means you keep ALL the flexibility without having to deal with
 any of the boilerplate code
 

diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst
@@ -361,6 +361,36 @@ iou (F)
 .. autofunction:: pytorch_lightning.metrics.functional.iou
     :noindex:
 
+mse (F)
+^^^^^^^
+
+.. autofunction:: pytorch_lightning.metrics.functional.mse
+    :noindex:
+
+rmse (F)
+^^^^^^^^
+
+.. autofunction:: pytorch_lightning.metrics.functional.rmse
+    :noindex:
+
+mae (F)
+^^^^^^^
+
+.. autofunction:: pytorch_lightning.metrics.functional.mae
+    :noindex:
+
+rmsle (F)
+^^^^^^^^^
+
+.. autofunction:: pytorch_lightning.metrics.functional.rmsle
+    :noindex:
+
+psnr (F)
+^^^^^^^^
+
+.. autofunction:: pytorch_lightning.metrics.functional.psnr
+    :noindex:
+
 stat_scores_multiple_classes (F)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

diff --git a/pytorch_lightning/core/hooks.py b/pytorch_lightning/core/hooks.py
@@ -39,7 +39,7 @@ def prepare_data(self):
                     # don't do this
                     self.something = else
 
-                def setup(step):
+                def setup(stage):
                     data = Load_data(...)
                     self.l1 = nn.Linear(28, data.num_classes)
 

diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -1321,7 +1321,7 @@ def prepare_data(self):
 
             model.prepare_data()
                 if ddp/tpu: init()
-            model.setup(step)
+            model.setup(stage)
             model.train_dataloader()
             model.val_dataloader()
             model.test_dataloader()

diff --git a/pytorch_lightning/metrics/__init__.py b/pytorch_lightning/metrics/__init__.py
@@ -1,11 +1,11 @@
 from pytorch_lightning.metrics.converters import numpy_metric, tensor_metric
 from pytorch_lightning.metrics.metric import Metric, TensorMetric, NumpyMetric
 from pytorch_lightning.metrics.regression import (
+    MAE,
     MSE,
+    PSNR,
     RMSE,
-    MAE,
-    RMSLE,
-    PSNR
+    RMSLE
 )
 from pytorch_lightning.metrics.classification import (
     Accuracy,
@@ -48,10 +48,10 @@
     'IoU',
 ]
 __regression_metrics = [
+    'MAE',
     'MSE',
+    'PSNR',
     'RMSE',
-    'MAE',
-    'RMSLE',
-    'PSNR'
+    'RMSLE'
 ]
 __all__ = __regression_metrics + __classification_metrics + ['SklearnMetric']
diff --git a/pytorch_lightning/metrics/functional/__init__.py b/pytorch_lightning/metrics/functional/__init__.py
@@ -20,3 +20,10 @@
     to_onehot,
     iou,
 )
+from pytorch_lightning.metrics.functional.regression import (
+    mae,
+    mse,
+    psnr,
+    rmse,
+    rmsle
+)
diff --git a/pytorch_lightning/metrics/functional/regression.py b/pytorch_lightning/metrics/functional/regression.py
@@ -4,6 +4,139 @@
 from pytorch_lightning.metrics.functional.reduction import reduce
 
 
+def mse(
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        reduction: str = 'elementwise_mean'
+) -> torch.Tensor:
+    """
+    Computes mean squared error
+
+    Args:
+        pred: estimated labels
+        target: ground truth labels
+        reduction: method for reducing mse (default: takes the mean)
+            Available reduction methods:
+
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum: add elements
+
+    Return:
+        Tensor with MSE
+
+    Example:
+
+        >>> x = torch.tensor([0., 1, 2, 3])
+        >>> y = torch.tensor([0., 1, 2, 2])
+        >>> mse(x, y)
+        tensor(0.2500)
+
+    """
+    mse = F.mse_loss(pred, target, reduction='none')
+    mse = reduce(mse, reduction=reduction)
+    return mse
+
+
+def rmse(
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        reduction: str = 'elementwise_mean'
+) -> torch.Tensor:
+    """
+    Computes root mean squared error
+
+    Args:
+        pred: estimated labels
+        target: ground truth labels
+        reduction: method for reducing rmse (default: takes the mean)
+            Available reduction methods:
+
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum: add elements
+
+    Return:
+        Tensor with RMSE
+
+
+        >>> x = torch.tensor([0., 1, 2, 3])
+        >>> y = torch.tensor([0., 1, 2, 2])
+        >>> rmse(x, y)
+        tensor(0.5000)
+
+    """
+    rmse = torch.sqrt(mse(pred, target, reduction=reduction))
+    return rmse
+
+
+def mae(
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        reduction: str = 'elementwise_mean'
+) -> torch.Tensor:
+    """
+    Computes mean absolute error
+
+    Args:
+        pred: estimated labels
+        target: ground truth labels
+        reduction: method for reducing mae (default: takes the mean)
+            Available reduction methods:
+
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum: add elements
+
+    Return:
+        Tensor with MAE
+
+    Example:
+
+        >>> x = torch.tensor([0., 1, 2, 3])
+        >>> y = torch.tensor([0., 1, 2, 2])
+        >>> mae(x, y)
+        tensor(0.2500)
+
+    """
+    mae = F.l1_loss(pred, target, reduction='none')
+    mae = reduce(mae, reduction=reduction)
+    return mae
+
+
+def rmsle(
+        pred: torch.Tensor,
+        target: torch.Tensor,
+        reduction: str = 'elementwise_mean'
+) -> torch.Tensor:
+    """
+    Computes root mean squared log error
+
+    Args:
+        pred: estimated labels
+        target: ground truth labels
+        reduction: method for reducing rmsle (default: takes the mean)
+            Available reduction methods:
+
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum: add elements
+
+    Return:
+        Tensor with RMSLE
+
+    Example:
+
+        >>> x = torch.tensor([0., 1, 2, 3])
+        >>> y = torch.tensor([0., 1, 2, 2])
+        >>> rmsle(x, y)
+        tensor(0.0207)
+
+    """
+    rmsle = mse(torch.log(pred + 1), torch.log(target + 1), reduction=reduction)
+    return rmsle
+
+
 def psnr(
     pred: torch.Tensor,
     target: torch.Tensor,
@@ -12,14 +145,22 @@ def psnr(
     reduction: str = 'elementwise_mean'
 ) -> torch.Tensor:
     """
-    Computes the peak signal-to-noise ratio metric
+    Computes the peak signal-to-noise ratio
 
     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).
+        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)
+            Available reduction methods:
+
+            - elementwise_mean: takes the mean
+            - none: pass array
+            - sum add elements
+
+    Return:
+        Tensor with PSNR score
 
     Example:
 
@@ -29,12 +170,15 @@ def psnr(
         >>> 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(float(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(torch.tensor(base)))
+
+    mse_score = mse(pred.view(-1), target.view(-1), reduction=reduction)
+    psnr_base_e = 2 * torch.log(data_range) - torch.log(mse_score)
+    psnr = psnr_base_e * (10 / torch.log(torch.tensor(base)))
+    return psnr