[metrics] change default behaviour of state dict (#4685)

* fix state dict

* Update docs/source/metrics.rst

Co-authored-by: Rohit Gupta <[email protected]>

* changelog

Co-authored-by: Rohit Gupta <[email protected]>
Co-authored-by: chaton <[email protected]>

CHANGELOG.md

@@ -15,6 +15,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
 
+- Metric states are no longer as default added to `state_dict` ([#4685](https://github.com/PyTorchLightning/pytorch-lightning/pull/))
+
+
 ### Deprecated
 
 

docs/source/metrics.rst

@@ -133,8 +133,8 @@ This metrics API is independent of PyTorch Lightning. Metrics can directly be us
 
 .. note::
 
-    Metric states will as default add their internal state to the models ``state_dict``.
-    To change this after initializing the metric the method ``.persistent(mode)`` can
+    Metric states are **not** added to the models ``state_dict`` by default.
+    To change this, after initializing the metric, the method ``.persistent(mode)`` can
     be used to enable (``mode=True``) or disable (``mode=False``) this behaviour.
 
 *********************

pytorch_lightning/metrics/metric.py

@@ -86,7 +86,7 @@ def __init__(
         self._reductions = {}
 
     def add_state(
-        self, name: str, default, dist_reduce_fx: Optional[Union[str, Callable]] = None, persistent: bool = True
+        self, name: str, default, dist_reduce_fx: Optional[Union[str, Callable]] = None, persistent: bool = False
     ):
         """
         Adds metric state variable. Only used by subclasses.
@@ -100,6 +100,7 @@ def add_state(
                 and ``torch.cat`` respectively, each with argument ``dim=0``. The user can also pass a custom
                 function in this parameter.
             persistent (Optional): whether the state will be saved as part of the modules ``state_dict``.
+                Default is ``False``.
 
         Note:
             Setting ``dist_reduce_fx`` to None will return the metric state synchronized across different processes.
@@ -278,7 +279,7 @@ def _apply(self, fn):
                                 f'or a list of torch.Tensor, but encountered {current_val}')
         return self
 
-    def persistent(self, mode: bool = True):
+    def persistent(self, mode: bool = False):
         """ Method for post-init to change if metric states should be saved to
             its state_dict
         """

tests/metrics/test_metric.py

@@ -1,6 +1,7 @@
 import pickle
-
+from collections import OrderedDict
 from distutils.version import LooseVersion
+
 import cloudpickle
 import numpy as np
 import pytest
@@ -167,3 +168,13 @@ def test_pickle(tmpdir):
     metric_loaded = cloudpickle.loads(metric_pickled)
 
     assert metric_loaded.compute() == 1
+
+
+def test_state_dict(tmpdir):
+    """ test that metric states can be removed and added to state dict """
+    metric = Dummy()
+    assert metric.state_dict() == OrderedDict()
+    metric.persistent(True)
+    assert metric.state_dict() == OrderedDict(x=0)
+    metric.persistent(False)
+    assert metric.state_dict() == OrderedDict()