Make metrics serializable

docs/source/engine.rst

@@ -69,7 +69,7 @@ Resuming the training
 It is possible to resume the training from a checkpoint and approximately reproduce original run's behaviour.
 Using Ignite, this can be easily done using :class:`~ignite.handlers.checkpoint.Checkpoint` handler. Engine provides two methods
 to serialize and deserialize its internal state :meth:`~ignite.engine.engine.Engine.state_dict` and
-:meth:`~ignite.engine.engine.Engine.load_state_dict`. In addition to serializing model, optimizer, lr scheduler etc user can
+:meth:`~ignite.engine.engine.Engine.load_state_dict`. In addition to serializing model, optimizer, lr scheduler, metrics, etc., user can
 store the trainer and then resume the training. For example:
 
 .. code-block:: python
@@ -82,8 +82,9 @@ store the trainer and then resume the training. For example:
     optimizer = ...
     lr_scheduler = ...
     data_loader = ...
+    metric = ...
 
-    to_save = {'trainer': trainer, 'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}
+    to_save = {'trainer': trainer, 'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler, 'metric': metric}
     handler = Checkpoint(to_save, DiskSaver('/tmp/training', create_dir=True))
     trainer.add_event_handler(Events.EPOCH_COMPLETED, handler)
     trainer.run(data_loader, max_epochs=100)
@@ -104,8 +105,9 @@ We can then restore the training from the last checkpoint.
     optimizer = ...
     lr_scheduler = ...
     data_loader = ...
+    metric = ...
 
-    to_load = {'trainer': trainer, 'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}
+    to_load = {'trainer': trainer, 'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler, 'metric': metric}
     checkpoint = torch.load(checkpoint_file)
     Checkpoint.load_objects(to_load=to_load, checkpoint=checkpoint)
 

ignite/handlers/checkpoint.py

@@ -710,7 +710,7 @@ def state_dict(self) -> "OrderedDict[str, List[Tuple[int, str]]]":
         return OrderedDict([("saved", [(p, f) for p, f in self._saved])])
 
     def load_state_dict(self, state_dict: Mapping) -> None:
-        """Method replace internal state of the class with provided state dict data.
+        """Method replaces internal state of the class with provided state dict data.
 
         Args:
             state_dict: a dict with "saved" key and list of ``(priority, filename)`` pairs as values.

ignite/metrics/accumulation.py

@@ -38,6 +38,7 @@ class VariableAccumulation(Metric):
     """
 
     required_output_keys = None
+    _state_dict_all_req_keys = ("accumulator", "num_examples")
 
     def __init__(
         self,

ignite/metrics/metric.py

@@ -1,12 +1,16 @@
 from abc import ABCMeta, abstractmethod
+from collections import OrderedDict
 from collections.abc import Mapping
 from functools import wraps
 from numbers import Number
-from typing import Any, Callable, cast, Dict, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, TYPE_CHECKING, Union
 
 import torch
+from torch import distributed as torch_dist
 
 import ignite.distributed as idist
+
+from ignite.base.mixins import Serializable
 from ignite.engine import CallableEventWithFilter, Engine, Events
 
 if TYPE_CHECKING:
@@ -216,7 +220,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         )
 
 
-class Metric(metaclass=ABCMeta):
+class Metric(Serializable, metaclass=ABCMeta):
     """
     Base class for all Metrics.
 
@@ -546,6 +550,41 @@ def is_attached(self, engine: Engine, usage: Union[str, MetricUsage] = EpochWise
         usage = self._check_usage(usage)
         return engine.has_event_handler(self.completed, usage.COMPLETED)
 
+    def state_dict(self) -> OrderedDict:
+        """Method returns state dict with attributes of the metric specified in its
+        `_state_dict_all_req_keys` attribute. Can be used to save internal state of the class.
+
+        If there's an active distributed configuration, some collective operations is done and
+        the list of values across ranks is saved under each attribute's name in the dict.
+        """
+        state = OrderedDict()
+        for attr_name in self._state_dict_all_req_keys:
+            attr = getattr(self, attr_name)
+            if idist.get_world_size() == 1:
+                state[attr_name] = [attr]
+            else:
+                if isinstance(attr, (float, torch.Tensor)):
+                    state[attr_name] = cast(List[Any], idist.all_gather(attr))
+                else:
+                    state[attr_name] = [None] * idist.get_world_size()
+                    torch_dist.all_gather_object(state[attr_name], attr)
+        return state
+
+    def load_state_dict(self, state_dict: Mapping) -> None:
+        """Method replaces internal state of the class with provided state dict data.
+
+        If there's an active distributed configuration, the process uses its rank to pick the proper value from
+        the list of values saved under each attribute's name in the dict.
+
+        Args:
+            state_dict: a dict containing attributes of the metric specified in its `_state_dict_all_req_keys`
+                attribute.
+        """
+        super().load_state_dict(state_dict)
+        rank = idist.get_local_rank()
+        for attr in self._state_dict_all_req_keys:
+            setattr(self, attr, state_dict[attr][rank])
+
     def __add__(self, other: Any) -> "MetricsLambda":
         from ignite.metrics.metrics_lambda import MetricsLambda
 

tests/ignite/metrics/test_metric.py

@@ -713,6 +713,7 @@ def test_distrib_nccl_gpu(distributed_context_single_node_nccl):
     _test_distrib_sync_all_reduce_decorator(device)
     _test_invalid_sync_all_reduce(device)
     _test_compute_with_sync_all_reduce_doesnt_change_attributes(device)
+    _test_distrib_state_dict(device)
 
 
 @pytest.mark.distributed
@@ -722,6 +723,7 @@ def test_distrib_gloo_cpu_or_gpu(distributed_context_single_node_gloo):
     _test_distrib_sync_all_reduce_decorator(device)
     _test_invalid_sync_all_reduce(device)
     _test_compute_with_sync_all_reduce_doesnt_change_attributes(device)
+    _test_distrib_state_dict(device)
 
 
 @pytest.mark.distributed
@@ -744,6 +746,7 @@ def test_multinode_distrib_gloo_cpu_or_gpu(distributed_context_multi_node_gloo):
     _test_distrib_sync_all_reduce_decorator(device)
     _test_invalid_sync_all_reduce(device)
     _test_compute_with_sync_all_reduce_doesnt_change_attributes(device)
+    _test_distrib_state_dict(device)
 
 
 @pytest.mark.multinode_distributed
@@ -1125,3 +1128,46 @@ def update(self, output):
 
     with pytest.raises(ValueError, match=r"Output should have 2 items of the same length"):
         engine.run([0] * 10)
+
+
+class DummyMetric4(Metric):
+    _state_dict_all_req_keys = ("number", "tensor", "object")
+
+    def __init__(self, value):
+        super().reset()
+        self.number = value
+        self.tensor = torch.tensor([value + 1])
+        self.object = {"a": [value + 2]}
+
+    def reset(self):
+        self.number = -1
+        self.tensor = torch.tensor([-2])
+        self.object = {"a": [-3]}
+
+    def update(output):
+        pass
+
+    def compute():
+        pass
+
+
+def test_state_dict():
+    metric = DummyMetric4(1)
+    state = metric.state_dict()
+    assert state.keys() == {"number", "tensor", "object"}
+    metric.reset()
+    metric.load_state_dict(state)
+    assert metric.number == 1
+    assert metric.tensor == torch.tensor([2])
+    assert metric.object == {"a": [3]}
+
+
+def _test_distrib_state_dict(device):
+    rank = idist.get_local_rank()
+    metric = DummyMetric4(rank)
+    state = metric.state_dict()
+    metric.reset()
+    metric.load_state_dict(state)
+    assert metric.number == rank
+    assert metric.tensor == torch.tensor([rank + 1])
+    assert metric.object == {"a": [rank + 2]}