Repository: Lightning-AI/pytorch-lightning
Pull request: Metric aggregation testing
<!--
Please note that we have freeze state on adding new features till v1.0 release.
Anyway, any new feature suggestion is welcome and you are free to draft a PR,
 but be aware that several refactoring will take place for this major release yielding in significant collision solving.
A recommendation for feature draft is draw just outline, do not implement all details and propagate the changes to codebase...
-->

## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.
-->

With PR #3245 we changed the way we are doing ddp sync and also introduced aggregation over multiple batches. This PR will add test for all metrics such that we are sure that they work in the following cases:
* Aggregation over multiple devices (aka running in ddp mode)
* Aggregation over multiple equal size batches
* Aggregation over multiple unequal size batches (important because a lot of assumptions that the final metric value can be computed as the mean over individual computed values break in this example)
* Aggregation over multiple devices and multiple batches
* Model integration

Should fix #3230

# Before submitting

- [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
- [ ] Did you make sure to update the documentation with your changes?
- [ ] Did you write any new necessary tests? 
- [ ] Did you verify new and existing tests pass locally with your changes?
- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## PR review    
Anyone in the community is free to review the PR once the tests have passed.     
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃


Problem statement:
Test all metrics against sklearn (with many input trials)
Hand-chosen values are not enough, we need to test with a large batch of inputs where possible. 
Something in this style, maybe with a fixed seed: 
```python
def test_auroc_versus_sklearn():
    for i in range(100):
        target = torch.randint(0, 2, size=(10, ))
        pred = torch.randint(0, 2, size=(10,))
        score_sk = sklearn.metrics.roc_auc_score(target.numpy(), pred.numpy())  # sklearn
        score_pl = auroc(pred, target)  # Pytorch Lightning 
        assert torch.allclose(torch.tensor(score_pl).float(), torch.tensor(score_sk).float())
```


Accuracy metric return tuple of length num_gpus on ddp in 0.9.1rc4
code:
```
        vid_acc = self.accuracy_video(video_labels_hat, video_labels)
        print(len(vid_acc), vid_acc)
        monitor = 0-vid_acc
```
In 0.9.1rc3, vid acc is a tensor, but in rc4, it changes to a tuple. I want to use `-vid_acc` as monitor, and I think it should be a tensor. 
Using rc4, in macos's cpu mode, it's a tensor. But in linux ddp mode, it's a tuple of length `num_gpus`.

Hints from issue discussion:
> maybe with a fixed seed

I don't think a seed is required here.
> > maybe with a fixed seed
> 
> I don't think a seed is required here.

I would agree here that all metric shall be deterministic so even using different seeds increase the coverage or cases and so far you pas the same values to the two functions it shall always give the same result, right?
> I would agree here that all metric shall be deterministic so even using different seeds increase the coverage or cases and so far you pas the same values to the two functions it shall always give the same result, right?

Agreed, this is one good point.
@awaelchli agree that we need such test. We already do that for many metrics:
https://github.com/PyTorchLightning/pytorch-lightning/blob/3910ad033074367f6abfe0001562db725a75cb73/tests/metrics/functional/test_classification.py#L38-L59
and
https://github.com/PyTorchLightning/pytorch-lightning/blob/3910ad033074367f6abfe0001562db725a75cb73/tests/metrics/functional/test_regression.py#L26-L44
tests are not deterministic, because no seed is used.

Commit message: fix pickle error

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b6f4ddf74662c..2d80e8b39bada 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -71,6 +71,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed determinism in `DDPSpawnBackend` when using `seed_everything` in main process ([#3335](https://github.com/PyTorchLightning/pytorch-lightning/pull/3335))
 
+- Fixed aggregation of metrics ([#3517](https://github.com/PyTorchLightning/pytorch-lightning/pull/3517))
+
 ## [0.9.0] - YYYY-MM-DD
 
 ### Added
diff --git a/pytorch_lightning/metrics/classification.py b/pytorch_lightning/metrics/classification.py
index a279cb3a4747e..bd1007059ef99 100644
--- a/pytorch_lightning/metrics/classification.py
+++ b/pytorch_lightning/metrics/classification.py
@@ -21,17 +21,18 @@
     auroc,
     average_precision,
     confusion_matrix,
+    _confmat_normalize,
     dice_score,
     f1_score,
     fbeta_score,
     iou,
     multiclass_precision_recall_curve,
     multiclass_roc,
-    precision,
     precision_recall_curve,
-    recall,
     roc,
+    precision_recall
 )
+from pytorch_lightning.metrics.functional.reduction import class_reduce
 from pytorch_lightning.metrics.metric import TensorMetric
 
 
@@ -44,8 +45,8 @@ class Accuracy(TensorMetric):
         >>> pred = torch.tensor([0, 1, 2, 3])
         >>> target = torch.tensor([0, 1, 2, 2])
         >>> metric = Accuracy()
-        >>> metric(pred, target).item()
-        0.75
+        >>> metric(pred, target)
+        tensor(0.7500)
 
     """
 
@@ -84,7 +85,14 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             A Tensor with the classification score.
         """
         return accuracy(pred=pred, target=target,
-                        num_classes=self.num_classes, class_reduction=self.class_reduction)
+                        num_classes=self.num_classes,
+                        class_reduction='none',
+                        return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        tps, sups = output['tps'], output['sups']
+        return class_reduce(tps, sups, sups, class_reduction=self.class_reduction)
 
 
 class ConfusionMatrix(TensorMetric):
@@ -135,16 +143,16 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             A Tensor with the confusion matrix.
         """
         return confusion_matrix(pred=pred, target=target,
-                                normalize=self.normalize,
+                                normalize=False,  # we normalize after ddp sync
                                 num_classes=self.num_classes)
 
-    def aggregate(self, *tensors: torch.Tensor) -> torch.Tensor:
-        """Aggregates results by stacking them instead of concatenating before averaging.
-
-        Returns:
-            the aggregated results
-        """
-        return torch.stack(tensors).mean(0)
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """ Confusion matrix normalization needs to happen after ddp sync """
+        confmat = output
+        if self.normalize:
+            confmat = _confmat_normalize(confmat)
+        return confmat
 
 
 class PrecisionRecallCurve(TensorMetric):
@@ -202,7 +210,8 @@ def forward(
             - recall values
             - threshold values
         """
-        return precision_recall_curve(pred=pred, target=target, sample_weight=sample_weight, pos_label=self.pos_label)
+        return precision_recall_curve(pred=pred, target=target,
+                                      sample_weight=sample_weight, pos_label=self.pos_label)
 
 
 class Precision(TensorMetric):
@@ -256,9 +265,15 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the classification score.
         """
-        return precision(pred=pred, target=target,
-                         num_classes=self.num_classes,
-                         class_reduction=self.class_reduction)
+        return precision_recall(pred=pred, target=target,
+                                num_classes=self.num_classes,
+                                class_reduction='none',
+                                return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        tps, fps, sups = output['tps'], output['fps'], output['sups']
+        return class_reduce(tps, tps + fps, sups, class_reduction=self.class_reduction)
 
 
 class Recall(TensorMetric):
@@ -313,10 +328,15 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the classification score.
         """
-        return recall(pred=pred,
-                      target=target,
-                      num_classes=self.num_classes,
-                      class_reduction=self.class_reduction)
+        return precision_recall(pred=pred, target=target,
+                                num_classes=self.num_classes,
+                                class_reduction='none',
+                                return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        tps, fns, sups = output['tps'], output['fns'], output['sups']
+        return class_reduce(tps, tps + fns, sups, class_reduction=self.class_reduction)
 
 
 class AveragePrecision(TensorMetric):
@@ -470,12 +490,28 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             torch.Tensor: classification score
         """
-        return fbeta_score(pred=pred, target=target,
-                           beta=self.beta, num_classes=self.num_classes,
-                           class_reduction=self.class_reduction)
+        return precision_recall(pred=pred, target=target,
+                                num_classes=self.num_classes,
+                                class_reduction='none',
+                                return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """ tps, fps, fns, sups needs to be synced before we do any calculations """
+        tps, fps, fns, sups = output['tps'], output['fps'], output['fns'], output['sups']
+
+        intermidiate_reduction = 'none' if self.class_reduction != "micro" else 'micro'
+        precision = class_reduce(tps, tps + fps, sups, class_reduction=intermidiate_reduction)
+        recall = class_reduce(tps, tps + fns, sups, class_reduction=intermidiate_reduction)
+
+        num = (1 + self.beta ** 2) * precision * recall
+        denom = ((self.beta ** 2) * precision + recall)
+        if intermidiate_reduction == 'micro':
+            return torch.sum(num) / torch.sum(denom)
+        return class_reduce(num, denom, sups, class_reduction=self.class_reduction)
 
 
-class F1(TensorMetric):
+class F1(FBeta):
     """
     Computes the F1 score, which is the harmonic mean of the precision and recall.
     It ranges between 1 and 0, where 1 is perfect and the worst value is 0.
@@ -507,7 +543,7 @@ def __init__(
 
             reduce_group: the process group to reduce metric results from DDP
         """
-        super().__init__(
+        super(TensorMetric, self).__init__(
             name="f1",
             reduce_group=reduce_group,
         )
@@ -515,21 +551,7 @@ def __init__(
         self.num_classes = num_classes
         assert class_reduction in ('micro', 'macro', 'weighted', 'none')
         self.class_reduction = class_reduction
-
-    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        """
-        Actual metric computation
-
-        Args:
-            pred: predicted labels
-            target: groundtruth labels
-
-        Return:
-            torch.Tensor: classification score
-        """
-        return f1_score(pred=pred, target=target,
-                        num_classes=self.num_classes,
-                        class_reduction=self.class_reduction)
+        self.beta = 1.0
 
 
 class ROC(TensorMetric):
diff --git a/pytorch_lightning/metrics/functional/classification.py b/pytorch_lightning/metrics/functional/classification.py
index 82e509684106d..5a02c546b9ee6 100644
--- a/pytorch_lightning/metrics/functional/classification.py
+++ b/pytorch_lightning/metrics/functional/classification.py
@@ -239,7 +239,8 @@ def accuracy(
         pred: torch.Tensor,
         target: torch.Tensor,
         num_classes: Optional[int] = None,
-        class_reduction: str = 'micro'
+        class_reduction: str = 'micro',
+        return_state=False
 ) -> torch.Tensor:
     """
     Computes the accuracy classification score
@@ -254,7 +255,8 @@ def accuracy(
             - ``'macro'``: calculate metrics for each label, and find their unweighted mean.
             - ``'weighted'``: calculate metrics for each label, and find their weighted mean.
             - ``'none'``: returns calculated metric per class
-
+        return_state: returns a internal state that can be ddp reduced
+            before doing the final calculation
     Return:
          A Tensor with the accuracy score.
 
@@ -271,10 +273,21 @@ def accuracy(
 
     tps, fps, tns, fns, sups = stat_scores_multiple_classes(
         pred=pred, target=target, num_classes=num_classes)
-
+    if return_state:
+        return {'tps': tps, 'sups': sups}
     return class_reduce(tps, sups, sups, class_reduction=class_reduction)
 
 
+def _confmat_normalize(cm):
+    """ Normalization function for confusion matrix """
+    cm = cm / cm.sum(-1, keepdim=True)
+    nan_elements = cm[torch.isnan(cm)].nelement()
+    if nan_elements != 0:
+        cm[torch.isnan(cm)] = 0
+        rank_zero_warn(f'{nan_elements} nan values found in confusion matrix have been replaced with zeros.')
+    return cm
+
+
 def confusion_matrix(
         pred: torch.Tensor,
         target: torch.Tensor,
@@ -312,11 +325,7 @@ def confusion_matrix(
     cm = bins.reshape(num_classes, num_classes).squeeze().float()
 
     if normalize:
-        cm = cm / cm.sum(-1, keepdim=True)
-        nan_elements = cm[torch.isnan(cm)].nelement()
-        if nan_elements != 0:
-            cm[torch.isnan(cm)] = 0
-            rank_zero_warn(f'{nan_elements} nan values found in confusion matrix have been replaced with zeros.')
+        cm = _confmat_normalize(cm)
 
     return cm
 
@@ -326,7 +335,8 @@ def precision_recall(
         target: torch.Tensor,
         num_classes: Optional[int] = None,
         class_reduction: str = 'micro',
-        return_support: bool = False
+        return_support: bool = False,
+        return_state: bool = False
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Computes precision and recall for different thresholds
@@ -343,6 +353,8 @@ def precision_recall(
             - ``'none'``: returns calculated metric per class
 
         return_support: returns the support for each class, need for fbeta/f1 calculations
+        return_state: returns a internal state that can be ddp reduced
+            before doing the final calculation
 
     Return:
         Tensor with precision and recall
@@ -359,6 +371,8 @@ def precision_recall(
 
     precision = class_reduce(tps, tps + fps, sups, class_reduction=class_reduction)
     recall = class_reduce(tps, tps + fns, sups, class_reduction=class_reduction)
+    if return_state:
+        return {'tps': tps, 'fps': fps, 'fns': fns, 'sups': sups}
     if return_support:
         return precision, recall, sups
     return precision, recall
diff --git a/pytorch_lightning/metrics/functional/regression.py b/pytorch_lightning/metrics/functional/regression.py
index 75d8f1adf9a86..a12123887c97f 100644
--- a/pytorch_lightning/metrics/functional/regression.py
+++ b/pytorch_lightning/metrics/functional/regression.py
@@ -9,7 +9,8 @@
 def mse(
         pred: torch.Tensor,
         target: torch.Tensor,
-        reduction: str = 'elementwise_mean'
+        reduction: str = 'elementwise_mean',
+        return_state: bool = False
 ) -> torch.Tensor:
     """
     Computes mean squared error
@@ -22,6 +23,8 @@ def mse(
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
             - ``'none'``: no reduction will be applied
+        return_state: returns a internal state that can be ddp reduced
+            before doing the final calculation
 
     Return:
         Tensor with MSE
@@ -35,6 +38,8 @@ def mse(
 
     """
     mse = F.mse_loss(pred, target, reduction='none')
+    if return_state:
+        return {'squared_error': mse.sum(), 'n_observations': torch.tensor(mse.numel())}
     mse = reduce(mse, reduction=reduction)
     return mse
 
@@ -42,7 +47,8 @@ def mse(
 def rmse(
         pred: torch.Tensor,
         target: torch.Tensor,
-        reduction: str = 'elementwise_mean'
+        reduction: str = 'elementwise_mean',
+        return_state=False
 ) -> torch.Tensor:
     """
     Computes root mean squared error
@@ -55,6 +61,8 @@ def rmse(
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
             - ``'none'``: no reduction will be applied
+        return_state: returns a internal state that can be ddp reduced
+            before doing the final calculation
 
     Return:
         Tensor with RMSE
@@ -66,14 +74,18 @@ def rmse(
         tensor(0.5000)
 
     """
-    rmse = torch.sqrt(mse(pred, target, reduction=reduction))
-    return rmse
+    mean_squared_error = mse(pred, target, reduction=reduction)
+    if return_state:
+        return {'squared_error': mean_squared_error.sum(),
+                'n_observations': torch.tensor(mean_squared_error.numel())}
+    return torch.sqrt(mean_squared_error)
 
 
 def mae(
         pred: torch.Tensor,
         target: torch.Tensor,
-        reduction: str = 'elementwise_mean'
+        reduction: str = 'elementwise_mean',
+        return_state: bool = False
 ) -> torch.Tensor:
     """
     Computes mean absolute error
@@ -86,6 +98,8 @@ def mae(
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
             - ``'none'``: no reduction will be applied
+        return_state: returns a internal state that can be ddp reduced
+            before doing the final calculation
 
     Return:
         Tensor with MAE
@@ -99,6 +113,8 @@ def mae(
 
     """
     mae = F.l1_loss(pred, target, reduction='none')
+    if return_state:
+        return {'absolute_error': mae.sum(), 'n_observations': torch.tensor(mae.numel())}
     mae = reduce(mae, reduction=reduction)
     return mae
 
@@ -140,7 +156,8 @@ def psnr(
     target: torch.Tensor,
     data_range: float = None,
     base: float = 10.0,
-    reduction: str = 'elementwise_mean'
+    reduction: str = 'elementwise_mean',
+    return_state: bool = False
 ) -> torch.Tensor:
     """
     Computes the peak signal-to-noise ratio
@@ -155,25 +172,30 @@ def psnr(
             - ``'elementwise_mean'``: takes the mean (default)
             - ``'sum'``: takes the sum
             - ``'none'``: no reduction will be applied
+        return_state: returns a internal state that can be ddp reduced
+            before doing the final calculation
 
     Return:
         Tensor with PSNR score
 
     Example:
 
-        >>> from pytorch_lightning.metrics.regression import PSNR
         >>> 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)
+        >>> psnr(pred, target)
         tensor(2.5527)
 
     """
     if data_range is None:
-        data_range = max(target.max() - target.min(), pred.max() - pred.min())
+        data_range = target.max() - target.min()
     else:
         data_range = torch.tensor(float(data_range))
 
+    if return_state:
+        return {'data_range': data_range,
+                'squared_error': F.mse_loss(pred, target, reduction='none').sum(),
+                'n_observations': torch.tensor(target.numel())}
+
     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)))
diff --git a/pytorch_lightning/metrics/metric.py b/pytorch_lightning/metrics/metric.py
index 45c50b084956f..e97f054f05f89 100644
--- a/pytorch_lightning/metrics/metric.py
+++ b/pytorch_lightning/metrics/metric.py
@@ -72,6 +72,7 @@ def __init__(self, name: str, reduce_group: Optional[Any] = None):
 
         self.reduce_group = reduce_group
 
+        # Buffer for holding aggregated state after each batch
         self._step_vals = []
 
         # Register hooks
@@ -131,9 +132,6 @@ def ddp_sync(self, tensor: Any):
 
         """
         gathered_tensors = apply_to_collection(tensor, torch.Tensor, gather_all_tensors_if_available, self.reduce_group)
-
-        self._step_vals.append(gathered_tensors)
-
         return gathered_tensors
 
     @staticmethod
@@ -150,7 +148,9 @@ def ddp_reduce(self, data: Any, output: Any):
 
         """
         synced = self.ddp_sync(output)
-        return self.aggregate(synced)
+        agg_val = self.aggregate(synced)
+        self._step_vals.append(agg_val)
+        return agg_val
 
     def aggregate(self, *tensors: torch.Tensor) -> torch.Tensor:
         """
@@ -163,17 +163,27 @@ def aggregate(self, *tensors: torch.Tensor) -> torch.Tensor:
             aggregated values
 
         """
-        try:
-            return torch.cat(tensors).mean(0)
-        except (ValueError, TypeError):
-            if isinstance(tensors[0], Mapping):
-                return {k: torch.stack([tensor[k] for tensor in tensors]).mean(0) for k in tensors[0].keys()}
-            elif isinstance(tensors[0], Sequence) and not isinstance(tensors[0], torch.Tensor):
-                return tuple([torch.stack(tmp).mean(0) for tmp in zip(*tensors)])
-            elif isinstance(tensors[0], torch.Tensor):
-                return torch.stack(tensors).mean(0)
-            else:
-                raise TypeError("unknown metric value format to aggregate")
+        # single tensor
+        if len(tensors) == 1:
+            tensors = tensors[0]
+            if isinstance(tensors, Mapping):
+                return {k: _stack_and_agg(tensors[k]) for k in tensors.keys()}
+            if isinstance(tensors, list):
+                return _stack_and_agg(tensors)
+            if isinstance(tensors, tuple):
+                return tensors
+            if isinstance(tensors, torch.Tensor):
+                return _stack_and_agg(tensors)
+
+        # multiple tensors (from aggregation over batches)
+        if isinstance(tensors[0], Mapping):
+            return {k: torch.stack([tensor[k] for tensor in tensors]).sum(0) for k in tensors[0].keys()}
+        if isinstance(tensors[0], Sequence):
+            return tuple([torch.stack(tmp).sum(0) for tmp in zip(*tensors)])
+        if isinstance(tensors[0], torch.Tensor):
+            return torch.stack(tensors).sum(0)
+
+        raise TypeError("unknown metric value format to aggregate")
 
     @staticmethod
     def compute(self, data: Any, output: Any):
@@ -192,7 +202,7 @@ def compute(self, data: Any, output: Any):
 
     @property
     def aggregated(self) -> torch.Tensor:
-        aggr = self.aggregate(*self._step_vals)
+        aggr = self.aggregate(*self._step_vals if len(self._step_vals) > 1 else self._step_vals)
         self.reset()
         return self.compute(self, None, aggr)
 
@@ -200,6 +210,13 @@ def reset(self):
         self._step_vals = []
 
 
+def _stack_and_agg(tensors):
+    """ Utility function for stacking and aggregating tensors """
+    if isinstance(tensors, list):
+        return torch.sum(torch.stack([t for t in tensors]), 0)
+    return tensors.squeeze() if tensors.numel() == 1 else tensors
+
+
 class TensorMetric(Metric):
     """
     Base class for metric implementation operating directly on tensors.
diff --git a/pytorch_lightning/metrics/regression.py b/pytorch_lightning/metrics/regression.py
index 66963f19bcc56..af8fcd3642ff4 100644
--- a/pytorch_lightning/metrics/regression.py
+++ b/pytorch_lightning/metrics/regression.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Sequence
+from typing import Sequence, Any
 
 import torch
 
@@ -67,7 +67,12 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the mse loss.
         """
-        return mse(pred, target, self.reduction)
+        return mse(pred, target, return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        sse, n = output['squared_error'], output['n_observations']
+        return sse / n
 
 
 class RMSE(Metric):
@@ -110,7 +115,13 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the rmse loss.
         """
-        return rmse(pred, target, self.reduction)
+        return rmse(pred, target, reduction='none', return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """ Squaring needs to happend after ddp sync"""
+        sse, n = output['squared_error'], output['n_observations']
+        return torch.sqrt(sse / n)
 
 
 class MAE(Metric):
@@ -153,7 +164,12 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the mae loss.
         """
-        return mae(pred, target, self.reduction)
+        return mae(pred, target, return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        sae, n = output['absolute_error'], output['n_observations']
+        return sae / n
 
 
 class RMSLE(Metric):
@@ -196,7 +212,14 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the rmsle loss.
         """
-        return rmsle(pred, target, self.reduction)
+        return mse(torch.log(pred + 1), torch.log(target + 1),
+                   self.reduction, return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """ Squaring needs to happend after ddp sync"""
+        sse, n = output['squared_error'], output['n_observations']
+        return torch.sqrt(sse / n)
 
 
 class PSNR(Metric):
@@ -245,7 +268,37 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with psnr score.
         """
-        return psnr(pred, target, self.data_range, self.base, self.reduction)
+        return psnr(pred, target, self.data_range, self.base, self.reduction, return_state=True)
+
+    def aggregate(self, *tensors: torch.Tensor) -> torch.Tensor:
+        """ Special aggregation function as the data range needs to be correct synced """
+        if len(tensors) == 1:
+            tensors = tensors[0]
+            output = {'data_range': torch.stack([t for t in tensors['data_range']]).max()}
+            output.update({k: torch.stack([t for t in tensors[k]]).sum(0) for k in tensors.keys() if k != 'data_range'})
+            return output
+
+        output = {'data_range': torch.stack([tensor['data_range'] for tensor in tensors]).max()}
+        output.update({k: torch.stack([tensor[k] for tensor in tensors]).sum(0) for k in tensors[0].keys() if k != 'data_range'})
+        return output
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """
+        Implement additionally metric computations to be done after the aggregation
+
+        Args:
+            data: input to forward method
+            output: output from the `aggregate` hook
+
+        Returns:
+            final metric value
+
+        """
+        sse, n, data_range = output['squared_error'], output['n_observations'], output['data_range']
+        psnr_base_e = 2 * torch.log(data_range) - torch.log(sse / n)
+        psnr = psnr_base_e * (10 / torch.log(torch.tensor(self.base)))
+        return psnr
 
 
 class SSIM(Metric):
diff --git a/pytorch_lightning/metrics/self_supervised.py b/pytorch_lightning/metrics/self_supervised.py
index b5e492662e854..9d97369dbf91e 100644
--- a/pytorch_lightning/metrics/self_supervised.py
+++ b/pytorch_lightning/metrics/self_supervised.py
@@ -18,6 +18,7 @@
 
 from pytorch_lightning.metrics.functional.self_supervised import embedding_similarity
 from pytorch_lightning.metrics.metric import TensorMetric
+from pytorch_lightning.utilities import rank_zero_warn
 
 
 class EmbeddingSimilarity(TensorMetric):
@@ -56,6 +57,9 @@ def __init__(
         assert reduction in ('none', 'sum', 'mean')
         self.reduction = reduction
 
+        rank_zero_warn('Please note that Metric `EmbeddingSimilarity` does not '
+                       'support aggregation.')
+
     def forward(self, batch: torch.Tensor) -> torch.Tensor:
         """
         Actual metric computation
@@ -71,3 +75,12 @@ def forward(self, batch: torch.Tensor) -> torch.Tensor:
                                     similarity=self.similarity,
                                     zero_diagonal=self.zero_diagonal,
                                     reduction=self.reduction)
+
+    @staticmethod
+    def ddp_reduce(self, data: Any, output: Any):
+        """ reduction for this metric does not make sense """
+        return output
+
+    @property
+    def aggregated(self):
+        raise ValueError('Metric `EmbeddingSimilarity` does not support aggregation.')
diff --git a/tests/base/model_test_epoch_ends.py b/tests/base/model_test_epoch_ends.py
index 429ac3401c6b5..ddce3f75b468b 100644
--- a/tests/base/model_test_epoch_ends.py
+++ b/tests/base/model_test_epoch_ends.py
@@ -2,6 +2,8 @@
 
 import torch
 
+from pytorch_lightning import EvalResult
+
 
 class TestEpochEndVariations(ABC):
 
@@ -73,3 +75,9 @@ def test_epoch_end__multiple_dataloaders(self, outputs):
         tqdm_dict = {'test_loss': test_loss_mean, 'test_acc': test_acc_mean}
         result = {'progress_bar': tqdm_dict}
         return result
+
+    def test_epoch_end__metrics(self, outputs):
+        aggregated_val = self.metric.aggregated
+        result = EvalResult()
+        result.log('metric_val', aggregated_val)
+        return result
diff --git a/tests/base/model_test_steps.py b/tests/base/model_test_steps.py
index 92b1d68e675a1..8fc473adee1b1 100644
--- a/tests/base/model_test_steps.py
+++ b/tests/base/model_test_steps.py
@@ -186,3 +186,10 @@ def test_step_result_preds(self, batch, batch_idx, optimizer_idx=None):
             result.write('vals', lst_of_dict, prediction_file)
 
         return result
+
+    def test_step__metrics(self, batch, batch_idx, **kwargs):
+        pred, target = batch
+        metric_val = self.metric(pred, target)
+        result = EvalResult()
+        result.log('metric_val', metric_val)
+        return result
diff --git a/tests/metrics/functional/test_regression.py b/tests/metrics/functional/test_regression.py
index bbba5a421ebfb..49a79f9424f13 100644
--- a/tests/metrics/functional/test_regression.py
+++ b/tests/metrics/functional/test_regression.py
@@ -90,7 +90,7 @@ def test_rmsle(pred, target, expected):
 ])
 def test_psnr_with_skimage(pred, target):
     score = psnr(pred=torch.tensor(pred),
-                 target=torch.tensor(target))
+                 target=torch.tensor(target), data_range=3)
     sk_score = ski_psnr(np.array(pred), np.array(target), data_range=3)
     assert torch.allclose(score, torch.tensor(sk_score, dtype=torch.float), atol=1e-3)
 
diff --git a/tests/metrics/test_aggregation.py b/tests/metrics/test_aggregation.py
new file mode 100644
index 0000000000000..8cdbfdd6ee050
--- /dev/null
+++ b/tests/metrics/test_aggregation.py
@@ -0,0 +1,296 @@
+import pytest
+import sys
+from collections import namedtuple
+from functools import partial
+import math
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import numpy as np
+
+from tests.base import EvalModelTemplate
+from pytorch_lightning import Trainer
+import tests.base.develop_utils as tutils
+from pytorch_lightning.metrics import (
+    Accuracy,
+    ConfusionMatrix,
+    PrecisionRecallCurve,
+    Precision,
+    Recall,
+    AveragePrecision,
+    AUROC,
+    FBeta,
+    F1,
+    ROC,
+    MulticlassROC,
+    MulticlassPrecisionRecallCurve,
+    DiceCoefficient,
+    IoU,
+    MAE,
+    MSE,
+    RMSE,
+    RMSLE,
+    PSNR,
+    SSIM,
+)
+
+from sklearn.metrics import (
+    accuracy_score,
+    confusion_matrix,
+    precision_recall_curve,
+    precision_score,
+    recall_score,
+    average_precision_score,
+    roc_auc_score,
+    fbeta_score,
+    f1_score,
+    roc_curve,
+    jaccard_score,
+    mean_squared_error,
+    mean_absolute_error,
+    mean_squared_log_error
+)
+
+from skimage.metrics import (
+    peak_signal_noise_ratio,
+    structural_similarity
+)
+
+# example structure
+TestCase = namedtuple('example', ['name', 'lightning_metric', 'comparing_metric', 'test_input'])
+
+# setup some standard testcases
+N_samples = 200
+multiclass_example = [(torch.randint(10, (N_samples,)), torch.randint(10, (N_samples,)))]
+binary_example = [(torch.randint(2, (N_samples,)), torch.randint(2, (N_samples,)))]
+multiclass_and_binary_example = [*multiclass_example, *binary_example]
+binary_example_logits = (torch.randint(2, (N_samples,)), torch.randint(5, (N_samples,)))
+multiclass_example_probs = (torch.randint(10, (N_samples,)), torch.randn((N_samples, 10)).softmax(-1))
+regression_example = [(torch.rand((N_samples,)), torch.rand((N_samples,)))]
+
+
+# construct additional test functions
+def root_mean_squared_error(x, y):
+    return math.sqrt(mean_squared_error(x, y))
+
+
+def root_mean_squared_log_error(x, y):
+    return math.sqrt(mean_squared_log_error(x, y))
+
+
+# Define testcases
+TESTS = [
+    TestCase('accuracy',
+             Accuracy,
+             accuracy_score,
+             multiclass_and_binary_example),
+    TestCase('confusion matrix without normalize',
+             ConfusionMatrix,
+             confusion_matrix,
+             multiclass_and_binary_example),
+    TestCase('confusion matrix with normalize',
+             partial(ConfusionMatrix, normalize=True),
+             partial(confusion_matrix, normalize='true'),
+             multiclass_and_binary_example),
+    # TestCase('precision recall curve',
+    #          PrecisionRecallCurve,
+    #          precision_recall_curve,
+    #          binary_example),
+    TestCase('precision',
+             Precision,
+             partial(precision_score, average='micro'),
+             multiclass_and_binary_example),
+    TestCase('recall',
+             Recall,
+             partial(recall_score, average='micro'),
+             multiclass_and_binary_example),
+    # TestCase('average_precision',
+    #          AveragePrecision,
+    #          average_precision_score,
+    #          binary_example),
+    # TestCase('auroc',
+    #          AUROC,
+    #          roc_auc_score,
+    #          binary_example),
+    TestCase('f beta',
+             partial(FBeta, beta=2),
+             partial(fbeta_score, average='micro', beta=2),
+             multiclass_and_binary_example),
+    TestCase('f1',
+             F1,
+             partial(f1_score, average='micro'),
+             multiclass_and_binary_example),
+    # TestCase('roc',
+    #          ROC,
+    #          roc_curve,
+    #          binary_example),
+    # TestCase('multiclass roc',
+    #          MulticlassROC,
+    #          multiclass_roc,
+    #          binary_example),
+    # TestCase('multiclass precision recall curve',
+    #          MulticlassPrecisionRecallCurve,
+    #          multiclass_precision_recall_curve,
+    #          binary_example),
+    # TestCase('dice coefficient',
+    #          DiceCoefficient,
+    #          partial(f1_score, average='micro'),
+    #          multiclass_and_binary_example),
+    # TestCase('intersection over union',
+    #          IoU,
+    #          partial(jaccard_score, average='macro'),
+    #          binary_example),
+    TestCase('mean squared error',
+             MSE,
+             mean_squared_error,
+             regression_example),
+    TestCase('root mean squared error',
+             RMSE,
+             root_mean_squared_error,
+             regression_example),
+    TestCase('mean absolute error',
+             MAE,
+             mean_absolute_error,
+             regression_example),
+    TestCase('root mean squared log error',
+             RMSLE,
+             root_mean_squared_log_error,
+             regression_example),
+    TestCase('peak signal-to-noise ratio',
+             partial(PSNR, data_range=10),
+             partial(peak_signal_noise_ratio, data_range=10),
+             regression_example),
+    # TestCase('structual similarity index measure',
+    #          SSIM,
+    #          structural_similarity,
+    #          regression_example)
+]
+
+
+# Utility test functions
+def idsfn(test):
+    """ Return id for current example being tested """
+    return test.name
+
+
+def _setup_ddp(rank, worldsize):
+    """ setup ddp enviroment for testing """
+    import os
+    os.environ['MASTER_ADDR'] = 'localhost'
+    # initialize the process group
+    dist.init_process_group("gloo", rank=rank, world_size=worldsize)
+
+
+def comparing_fn(lightning_val, comparing_val, rtol=1e-03, atol=1e-08):
+    """ function for comparing output, both multi and single output"""
+    # multi output
+    if isinstance(comparing_val, tuple):
+        for l_score, c_score in zip(lightning_val, comparing_val):
+            assert np.allclose(l_score.numpy(), c_score, rtol, atol)
+    else:  # single output
+        assert np.allclose(lightning_val.numpy(), comparing_val, rtol, atol)
+
+
+# ===== Tests start here =====
+def _test_ddp_single_batch(rank, worldsize, lightning_metric, comparing_metric, test_inputs):
+    """ ddp testing function, divide test_inputs equally between all processes """
+    _setup_ddp(rank, worldsize)
+
+    # Setup metric for ddp
+    lightning_metric = lightning_metric()
+    for test_input in test_inputs:
+        # rank 0 receives sample 0,2,4,...
+        # rank 1 receives sample 1,3,5,...
+        lightning_val = lightning_metric(*[ti[rank::2] for ti in test_input])
+
+        comparing_val = comparing_metric(*[ti.numpy() for ti in reversed(test_input)])
+
+        comparing_fn(lightning_val, comparing_val)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="DDP not available on windows")
+@pytest.mark.parametrize("test", TESTS, ids=idsfn)
+def test_ddp(test):
+    """Make sure that metrics are correctly sync and reduced in DDP mode"""
+    tutils.reset_seed()
+    tutils.set_random_master_port()
+
+    worldsize = 2
+    mp.spawn(_test_ddp_single_batch,
+             args=(worldsize,
+                   test.lightning_metric,
+                   test.comparing_metric,
+                   test.test_input),
+             nprocs=worldsize)
+
+
+@pytest.mark.parametrize("test", TESTS, ids=idsfn)
+def test_multi_batch(test):
+    """ test that aggregation works for multiple batches """
+    lightning_metric = test.lightning_metric()
+    comparing_metric = test.comparing_metric
+
+    for test_input in test.test_input:
+        for i in range(2):  # for lightning device in 2 artificially batches
+            # first batch consist of samples 0,2,4,...
+            # second batch consist of samples 1,3,5,...
+            _ = lightning_metric(*[ti[i::2] for ti in test_input])
+        lightning_val = lightning_metric.aggregated
+        comparing_val = comparing_metric(*[ti.numpy() for ti in reversed(test_input)])
+
+        comparing_fn(lightning_val, comparing_val)
+
+
+@pytest.mark.parametrize("test", TESTS, ids=idsfn)
+def test_multi_batch_unequal_sizes(test):
+    """ test that aggregation works for multiple batches with uneven sizes """
+    lightning_metric = test.lightning_metric()
+    comparing_metric = test.comparing_metric
+
+    for test_input in test.test_input:
+        for i in range(2):  # for lightning device in 2 artificially batches
+            if i == 0:  # allocate 3/4 of data to the first batch
+                _ = lightning_metric(*[ti[:int(3 / 4 * len(ti))] for ti in test_input])
+            else:
+                _ = lightning_metric(*[ti[int(3 / 4 * len(ti)):] for ti in test_input])
+        lightning_val = lightning_metric.aggregated
+        comparing_val = comparing_metric(*[ti.numpy() for ti in reversed(test_input)])
+
+        comparing_fn(lightning_val, comparing_val)
+
+
+def _test_ddp_multi_batch(rank, worldsize, lightning_metric, comparing_metric, test_inputs):
+    """ ddp testing function, test that metric works with aggregation over multiple
+        devices and multiple batches """
+    _setup_ddp(rank, worldsize)
+
+    # Setup metric for ddp
+    lightning_metric = lightning_metric()
+    for test_input in test_inputs:
+        for i in range(2):  # artificially divide samples between batches and processes
+            # rank 0, batch 0 consist of samples 0,4,8,...
+            # rank 0, batch 1 consist of samples 1,5,9,...
+            # rank 1, batch 0 consist of samples 2,6,10,...
+            # rank 1, batch 1 consist of samples 3,7,11,...
+            _ = lightning_metric(*[ti[i + worldsize * rank::4] for ti in test_input])
+        lightning_val = lightning_metric.aggregated
+        comparing_val = comparing_metric(*[ti.numpy() for ti in reversed(test_input)])
+
+        comparing_fn(lightning_val, comparing_val)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="DDP not available on windows")
+@pytest.mark.parametrize("test", TESTS, ids=idsfn)
+def test_ddp_multi_batch(test):
+    """ test that aggregation works fine with in DDP mode and multiple batches """
+    tutils.reset_seed()
+    tutils.set_random_master_port()
+
+    worldsize = 2
+    mp.spawn(_test_ddp_multi_batch,
+             args=(worldsize,
+                   test.lightning_metric,
+                   test.comparing_metric,
+                   test.test_input),
+             nprocs=worldsize)
diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_metrics.py
index 5a8589f8f254a..f3395a551ef62 100644
--- a/tests/metrics/test_metrics.py
+++ b/tests/metrics/test_metrics.py
@@ -190,7 +190,9 @@ def test_saving_pickable(tmpdir, metric: Metric):
     assert results_before_save == results_after_load
 
 
-def check_call_order():
+def test_correct_call_order():
+    """ Check that hooks are called in the expected order """
+
     class DummyMetric(Metric):
         def __init__(self):
             super().__init__("dummy")
@@ -249,6 +251,7 @@ def compute(self, data: Any, output: Any):
         "ddp_reduce",
         "ddp_sync",
         "aggregate",
+        "compute"
     ]
     aggr = metric.aggregated
     assert metric.call_history == [
@@ -259,9 +262,11 @@ def compute(self, data: Any, output: Any):
         "ddp_reduce",
         "ddp_sync",
         "aggregate",
+        "compute",
         "aggregated",
         "aggregate",
         "reset",
+        "compute"
     ]
     assert torch.allclose(aggr, result)
     _ = metric(torch.tensor(2.0), torch.tensor(1.0))
@@ -273,15 +278,18 @@ def compute(self, data: Any, output: Any):
         "ddp_reduce",
         "ddp_sync",
         "aggregate",
+        "compute",
         "aggregated",
         "aggregate",
         "reset",
+        "compute",
         "input_convert",
         "forward",
         "output_convert",
         "ddp_reduce",
         "ddp_sync",
         "aggregate",
+        "compute"
     ]
 
     metric = DummyMetric()
@@ -290,7 +298,7 @@ def compute(self, data: Any, output: Any):
 
     aggregated = metric.aggregated
 
-    assert torch.allclose(aggregated, torch.tensor(2.0))
+    assert torch.allclose(aggregated, torch.tensor(4.0))
 
     assert metric.call_history == [
         "init",
@@ -300,13 +308,16 @@ def compute(self, data: Any, output: Any):
         "ddp_reduce",
         "ddp_sync",
         "aggregate",
+        "compute",
         "input_convert",
         "forward",
         "output_convert",
         "ddp_reduce",
         "ddp_sync",
         "aggregate",
+        "compute",
         "aggregated",
         "aggregate",
         "reset",
+        "compute",
     ]


Code review comments:
diff_hunk: @@ -35,14 +38,17 @@ def mse(
 
     """
     mse = F.mse_loss(pred, target, reduction='none')
+    if return_state:
+        return {'squared_error': mse.sum(), 'n_observations': torch.tensor(mse.numel())}
     mse = reduce(mse, reduction=reduction)
     return mse
 
 
 def rmse(
         pred: torch.Tensor,
         target: torch.Tensor,
-        reduction: str = 'elementwise_mean'
+        reduction: str = 'elementwise_mean',
+        return_state=False
path: pytorch_lightning/metrics/functional/regression.py
text: @user1:
```suggestion
        return_state: bool = False
```

diff_hunk: @@ -110,7 +115,13 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the rmse loss.
         """
-        return rmse(pred, target, self.reduction)
+        return rmse(pred, target, reduction='none', return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """ Squaring needs to happend after ddp sync"""
path: pytorch_lightning/metrics/regression.py
text: @user1:
```suggestion
        """ Squaring needs to happend after ddp sync """
```

diff_hunk: @@ -507,29 +543,15 @@ def __init__(
 
             reduce_group: the process group to reduce metric results from DDP
         """
-        super().__init__(
+        super(TensorMetric, self).__init__(
path: pytorch_lightning/metrics/classification.py
text: @user1:
can't we just call fbeta init here with beta=1 and then overwrite the name? should be less code, no?

@author:
I have changed it to your suggestion 👍

diff_hunk: @@ -239,7 +239,8 @@ def accuracy(
         pred: torch.Tensor,
         target: torch.Tensor,
         num_classes: Optional[int] = None,
-        class_reduction: str = 'micro'
+        class_reduction: str = 'micro',
+        return_state=False
path: pytorch_lightning/metrics/functional/classification.py
text: @user1:
```suggestion
        return_state: bool = False
```

diff_hunk: @@ -245,7 +268,37 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with psnr score.
         """
-        return psnr(pred, target, self.data_range, self.base, self.reduction)
+        return psnr(pred, target, self.data_range, self.base, self.reduction, return_state=True)
+
+    def aggregate(self, *tensors: torch.Tensor) -> torch.Tensor:
+        """ Special aggregation function as the data range needs to be correct synced """
path: pytorch_lightning/metrics/regression.py
text: @user1:
```suggestion
        """ Special aggregation function as the data range needs to be correctly synced """
```

diff_hunk: @@ -196,7 +212,14 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Return:
             A Tensor with the rmsle loss.
         """
-        return rmsle(pred, target, self.reduction)
+        return mse(torch.log(pred + 1), torch.log(target + 1),
+                   self.reduction, return_state=True)
+
+    @staticmethod
+    def compute(self, data: Any, output: Any):
+        """ Squaring needs to happend after ddp sync"""
path: pytorch_lightning/metrics/regression.py
text: @user1:
```suggestion
        """ Squaring needs to happend after ddp sync """
```