[Metrics] Allow for running accumulation computations

CHANGELOG.md

@@ -21,6 +21,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `R2Score` metric ([#5241](https://github.com/PyTorchLightning/pytorch-lightning/pull/5241))
 
+- Added metrics' arguments allowing for running accumulation computations ([#5193](https://github.com/PyTorchLightning/pytorch-lightning/pull/5193))
+
 
 ### Changed
 

pytorch_lightning/metrics/classification/accuracy.py

@@ -77,6 +77,9 @@ class Accuracy(Metric):
         dist_sync_fn:
             Callback that performs the allgather operation on the metric state. When ``None``, DDP
             will be used to perform the allgather
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example:
 
@@ -104,12 +107,14 @@ def __init__(
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
         dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
             dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
 
         self.add_state("correct", default=torch.tensor(0), dist_reduce_fx="sum")

pytorch_lightning/metrics/classification/average_precision.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional, Any, Union, List
+from typing import Any, Callable, List, Optional, Union
 
 import torch
 
@@ -50,6 +50,12 @@ class AveragePrecision(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example (binary case):
 
@@ -78,11 +84,15 @@ def __init__(
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
 
         self.num_classes = num_classes

pytorch_lightning/metrics/classification/confusion_matrix.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import torch
 
@@ -61,6 +61,12 @@ class ConfusionMatrix(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example:
 
@@ -81,12 +87,16 @@ def __init__(
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
 
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
         self.num_classes = num_classes
         self.normalize = normalize

pytorch_lightning/metrics/classification/f_beta.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import torch
 
@@ -66,6 +66,12 @@ class FBeta(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example:
 
@@ -88,9 +94,15 @@ def __init__(
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
-            compute_on_step=compute_on_step, dist_sync_on_step=dist_sync_on_step, process_group=process_group,
+            compute_on_step=compute_on_step,
+            dist_sync_on_step=dist_sync_on_step,
+            process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
 
         self.num_classes = num_classes

pytorch_lightning/metrics/classification/precision_recall.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import torch
 
@@ -58,6 +58,12 @@ class Precision(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example:
 
@@ -78,11 +84,15 @@ def __init__(
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
 
         self.num_classes = num_classes

pytorch_lightning/metrics/classification/precision_recall_curve.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional, Any, Union, Tuple, List
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 import torch
 
@@ -50,6 +50,12 @@ class PrecisionRecallCurve(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example (binary case):
 
@@ -89,11 +95,15 @@ def __init__(
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
 
         self.num_classes = num_classes

pytorch_lightning/metrics/classification/roc.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional, Any, Union, List, Tuple
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 import torch
 
@@ -50,6 +50,12 @@ class ROC(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example (binary case):
 
@@ -91,11 +97,15 @@ def __init__(
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
+        dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
 
         self.num_classes = num_classes

pytorch_lightning/metrics/metric.py

@@ -22,7 +22,7 @@
 
 from pytorch_lightning.metrics.utils import _flatten, dim_zero_cat, dim_zero_mean, dim_zero_sum
 from pytorch_lightning.utilities.apply_func import apply_to_collection
-from pytorch_lightning.utilities.distributed import gather_all_tensors
+from pytorch_lightning.utilities.distributed import gather_all_tensors, all_except_rank_zero
 
 
 class Metric(nn.Module, ABC):
@@ -56,26 +56,32 @@ class Metric(nn.Module, ABC):
         dist_sync_fn:
             Callback that performs the allgather operation on the metric state. When `None`, DDP
             will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
     """
     def __init__(
         self,
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
         dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__()
 
         self.dist_sync_on_step = dist_sync_on_step
         self.compute_on_step = compute_on_step
         self.process_group = process_group
         self.dist_sync_fn = dist_sync_fn
+        self.auto_reset_on_compute = auto_reset_on_compute
         self._to_sync = True
 
         self.update = self._wrap_update(self.update)
         self.compute = self._wrap_compute(self.compute)
         self._computed = None
         self._forward_cache = None
+        self._internal_reset = False
 
         # initialize state
         self._defaults = {}
@@ -217,6 +223,7 @@ def wrapped_func(*args, **kwargs):
                 self._sync_dist(dist_sync_fn)
 
             self._computed = compute(*args, **kwargs)
+            self._internal_reset = self._to_sync
             self.reset()
 
             return self._computed
@@ -238,17 +245,26 @@ def compute(self):  # pylint: disable=E0202
         """
         pass
 
-    def reset(self):
-        """
-        This method automatically resets the metric state variables to their default value.
-        """
+    def _reset_to_default(self):
+        """ reset metric state to their default values """
         for attr, default in self._defaults.items():
             current_val = getattr(self, attr)
             if isinstance(current_val, torch.Tensor):
                 setattr(self, attr, deepcopy(default).to(current_val.device))
             else:
                 setattr(self, attr, deepcopy(default))
 
+    def reset(self):
+        """
+        This method automatically resets the metric state variables to their default value.
+        """
+        if self.auto_reset_on_compute or not self._internal_reset:
+            reset_fn = self._reset_to_default
+        else:
+            reset_fn = all_except_rank_zero(self._reset_to_default)
+        reset_fn()
+        self._internal_reset = False
+
     def __getstate__(self):
         # ignore update and compute functions for pickling
         return {k: v for k, v in self.__dict__.items() if k not in ["update", "compute"]}

pytorch_lightning/metrics/regression/explained_variance.py

@@ -57,6 +57,12 @@ class ExplainedVariance(Metric):
             before returning the value at the step. default: False
         process_group:
             Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather. default: None
+        auto_reset_on_compute:
+            Specify if ``reset()`` should be called automatically after each ``compute()``.
+            Disabling it allows for calculate running accumulated metrics. default: True
 
     Example:
 
@@ -81,12 +87,14 @@ def __init__(
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
         dist_sync_fn: Callable = None,
+        auto_reset_on_compute: bool = True
     ):
         super().__init__(
             compute_on_step=compute_on_step,
             dist_sync_on_step=dist_sync_on_step,
             process_group=process_group,
             dist_sync_fn=dist_sync_fn,
+            auto_reset_on_compute=auto_reset_on_compute
         )
         allowed_multioutput = ('raw_values', 'uniform_average', 'variance_weighted')
         if multioutput not in allowed_multioutput: