Added state_dict support for MetricsLambda and contrib metrics

vfdev-5 · vfdev-5 · commit 0cd005416f31 · 2023-10-16T12:00:45.000+02:00
diff --git a/ignite/contrib/metrics/regression/canberra_metric.py b/ignite/contrib/metrics/regression/canberra_metric.py
@@ -63,6 +63,7 @@ class CanberraMetric(_BaseRegression):
         - Fixed implementation: ``abs`` in denominator.
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors",)
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/fractional_absolute_error.py b/ignite/contrib/metrics/regression/fractional_absolute_error.py
@@ -58,6 +58,7 @@ class FractionalAbsoluteError(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors", "_num_examples")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/fractional_bias.py b/ignite/contrib/metrics/regression/fractional_bias.py
@@ -58,6 +58,7 @@ class FractionalBias(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors", "_num_examples")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/geometric_mean_absolute_error.py b/ignite/contrib/metrics/regression/geometric_mean_absolute_error.py
@@ -58,6 +58,7 @@ class GeometricMeanAbsoluteError(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors", "_num_examples")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/geometric_mean_relative_absolute_error.py b/ignite/contrib/metrics/regression/geometric_mean_relative_absolute_error.py
@@ -69,6 +69,7 @@ class GeometricMeanRelativeAbsoluteError(_BaseRegression):
 
             0.0...
     """
+    _state_dict_all_req_keys = ("_predictions", "_targets")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/manhattan_distance.py b/ignite/contrib/metrics/regression/manhattan_distance.py
@@ -59,6 +59,7 @@ class ManhattanDistance(_BaseRegression):
         - Fixed sklearn compatibility.
         - Workes with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors",)
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/maximum_absolute_error.py b/ignite/contrib/metrics/regression/maximum_absolute_error.py
@@ -58,6 +58,7 @@ class MaximumAbsoluteError(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_max_of_absolute_errors",)
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/mean_absolute_relative_error.py b/ignite/contrib/metrics/regression/mean_absolute_relative_error.py
@@ -58,6 +58,7 @@ class MeanAbsoluteRelativeError(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_absolute_relative_errors", "_num_samples")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/mean_error.py b/ignite/contrib/metrics/regression/mean_error.py
@@ -55,6 +55,7 @@ class MeanError(_BaseRegression):
 
             0.625...
     """
+    _state_dict_all_req_keys = ("_sum_of_errors", "_num_examples")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/mean_normalized_bias.py b/ignite/contrib/metrics/regression/mean_normalized_bias.py
@@ -58,6 +58,7 @@ class MeanNormalizedBias(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors", "_num_examples")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/r2_score.py b/ignite/contrib/metrics/regression/r2_score.py
@@ -56,6 +56,7 @@ class R2Score(_BaseRegression):
     .. versionchanged:: 0.4.3
         Works with DDP.
     """
+    _state_dict_all_req_keys = ("_num_examples", "_sum_of_errors", "_y_sq_sum", "_y_sum")
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/contrib/metrics/regression/wave_hedges_distance.py b/ignite/contrib/metrics/regression/wave_hedges_distance.py
@@ -57,6 +57,7 @@ class WaveHedgesDistance(_BaseRegression):
     .. versionchanged:: 0.4.5
         - Works with DDP.
     """
+    _state_dict_all_req_keys = ("_sum_of_errors",)
 
     @reinit__is_reduced
     def reset(self) -> None:
diff --git a/ignite/metrics/classification_report.py b/ignite/metrics/classification_report.py
@@ -131,12 +131,12 @@ def _wrapper(
             dict_obj[_get_label_for_class(idx)] = {
                 "precision": p_label.item(),
                 "recall": re[idx].item(),
-                "f{0}-score".format(beta): f[idx].item(),
+                f"f{beta}-score": f[idx].item(),
             }
         dict_obj["macro avg"] = {
             "precision": a_pr.item(),
             "recall": a_re.item(),
-            "f{0}-score".format(beta): a_f.item(),
+            f"f{beta}-score": a_f.item(),
         }
         return dict_obj if output_dict else json.dumps(dict_obj)
 
diff --git a/ignite/metrics/metrics_lambda.py b/ignite/metrics/metrics_lambda.py
@@ -1,4 +1,6 @@
 import itertools
+from collections import OrderedDict
+from collections.abc import Mapping
 from typing import Any, Callable, Optional, Union
 
 import torch
@@ -90,9 +92,11 @@ def Fbeta(r, p, beta):
             assert not precision.is_attached(engine)
     """
 
+    _state_dict_all_req_keys = ("_updated", "args", "kwargs")
+
     def __init__(self, f: Callable, *args: Any, **kwargs: Any) -> None:
         self.function = f
-        self.args = args
+        self.args = list(args)  # we need args to be a list instead of a tuple for state_dict/load_state_dict feature
         self.kwargs = kwargs
         self.engine: Optional[Engine] = None
         self._updated = False
diff --git a/ignite/metrics/precision.py b/ignite/metrics/precision.py
@@ -13,7 +13,7 @@
 
 
 class _BasePrecisionRecall(_BaseClassification):
-    _state_dict_all_req_keys = ("_numerator", "_denominator", "_weight")
+    _state_dict_all_req_keys = ("_numerator", "_denominator", "_weight", "_updated", "eps")
 
     def __init__(
         self,
diff --git a/ignite/utils.py b/ignite/utils.py
@@ -6,7 +6,7 @@
 import shutil
 import warnings
 from pathlib import Path
-from typing import Any, Callable, cast, Dict, Optional, TextIO, Tuple, Type, TypeVar, Union
+from typing import Any, Callable, cast, Dict, List, Optional, TextIO, Tuple, Type, TypeVar, Union
 
 import torch
 
@@ -93,7 +93,14 @@ def _tree_map(
 class _CollectionItem:
     types_as_collection_item: Tuple = (int, float, torch.Tensor)
 
-    def __init__(self, collection: Union[collections.Mapping, collections.Sequence], key: Union[int, str]) -> None:
+    def __init__(self, collection: Union[Dict, List], key: Union[int, str]) -> None:
+        if not isinstance(collection, (dict, list)):
+            raise TypeError(
+                f"Input type is expected to be a mapping or list, but got {type(collection)} " f"for input key '{key}'."
+            )
+        if isinstance(collection, list) and isinstance(key, str):
+            raise ValueError("Key should be int for collection of type list")
+
         self.collection = collection
         self.key = key
 
@@ -105,7 +112,7 @@ def value(self) -> Any:
 
     @staticmethod
     def wrap(
-        object: Union[collections.Mapping, collections.Sequence], key: Union[int, str], value: Any
+        object: Union[Dict, List], key: Union[int, str], value: Any
     ) -> Union[Any, "_CollectionItem"]:
         return (
             _CollectionItem(object, key)
@@ -116,15 +123,15 @@ def wrap(
 
 def _tree_apply2(
     func: Callable,
-    x: Union[Any, collections.Sequence, collections.Mapping],
+    x: Union[Any, List, Dict],
     y: Union[Any, collections.Sequence, collections.Mapping],
 ) -> None:
-    if isinstance(x, collections.Mapping) and isinstance(y, collections.Mapping):
+    if isinstance(x, dict) and isinstance(y, collections.Mapping):
         for k, v in x.items():
             if k not in y:
                 raise ValueError(f"Key '{k}' from x is not found in y: {y.keys()}")
             _tree_apply2(func, _CollectionItem.wrap(x, k, v), y[k])
-    elif isinstance(x, collections.Sequence) and isinstance(y, collections.Sequence):
+    elif isinstance(x, list) and isinstance(y, collections.Sequence):
         if len(x) != len(y):
             raise ValueError(f"Size of y: {len(y)} does not match the size of x: '{len(x)}'")
         for i, (v1, v2) in enumerate(zip(x, y)):
diff --git a/tests/ignite/metrics/test_accumulation.py b/tests/ignite/metrics/test_accumulation.py
@@ -152,6 +152,16 @@ def update_fn(engine, batch):
         np.array(state.metrics["agg_custom_var"]), true_result_fn(custom_variable), decimal=5
     )
 
+    metric_state = custom_var_mean.state_dict()
+    saved_num_examples = custom_var_mean.num_examples
+    saved_accumulator = custom_var_mean.accumulator
+    custom_var_mean.reset()
+    assert custom_var_mean.num_examples == 0
+    assert custom_var_mean.accumulator == 0
+    custom_var_mean.load_state_dict(metric_state)
+    assert custom_var_mean.num_examples == saved_num_examples
+    assert (custom_var_mean.accumulator == saved_accumulator).all()
+
 
 def test_compute_mean_std():
     n = 8
diff --git a/tests/ignite/metrics/test_accuracy.py b/tests/ignite/metrics/test_accuracy.py
@@ -405,6 +405,14 @@ def update(engine, i):
 
         assert pytest.approx(res) == true_res
 
+        metric_state = acc.state_dict()
+        saved__num_correct = acc._num_correct
+        saved__num_examples = acc._num_examples
+        acc.reset()
+        acc.load_state_dict(metric_state)
+        assert acc._num_examples == saved__num_examples
+        assert (acc._num_correct == saved__num_correct).all()
+
     metric_devices = ["cpu"]
     if device.type != "xla":
         metric_devices.append(idist.device())
diff --git a/tests/ignite/metrics/test_classification_report.py b/tests/ignite/metrics/test_classification_report.py
@@ -60,6 +60,22 @@ def update(engine, i):
         assert sklearn_result["macro avg"]["recall"] == pytest.approx(res["macro avg"]["recall"])
         assert sklearn_result["macro avg"]["f1-score"] == pytest.approx(res["macro avg"]["f1-score"])
 
+        metric_state = classification_report.state_dict()
+        classification_report.reset()
+        classification_report.load_state_dict(metric_state)
+        res2 = classification_report.compute()
+        if not output_dict:
+            res2 = json.loads(res2)
+
+        for i in range(n_classes):
+            label_i = labels[i] if labels else str(i)
+            assert res2[label_i]["precision"] == res[label_i]["precision"]
+            assert res2[label_i]["f1-score"] == res[label_i]["f1-score"]
+            assert res2[label_i]["recall"] == res[label_i]["recall"]
+        assert res2["macro avg"]["precision"] == res["macro avg"]["precision"]
+        assert res2["macro avg"]["recall"] == res["macro avg"]["recall"]
+        assert res2["macro avg"]["f1-score"] == res["macro avg"]["f1-score"]
+
     for i in range(5):
         torch.manual_seed(12 + rank + i)
         # check multiple random inputs as random exact occurencies are rare
@@ -121,7 +137,6 @@ def update(engine, i):
         sklearn_result = sklearn_classification_report(np_y_true, np_y_preds, output_dict=True, zero_division=1)
 
         for i in range(n_classes):
-            torch.manual_seed(12 + rank + i)
             label_i = labels[i] if labels else str(i)
             assert sklearn_result[str(i)]["precision"] == pytest.approx(res[label_i]["precision"])
             assert sklearn_result[str(i)]["f1-score"] == pytest.approx(res[label_i]["f1-score"])
@@ -130,7 +145,8 @@ def update(engine, i):
         assert sklearn_result["macro avg"]["recall"] == pytest.approx(res["macro avg"]["recall"])
         assert sklearn_result["macro avg"]["f1-score"] == pytest.approx(res["macro avg"]["f1-score"])
 
-    for _ in range(3):
+    for i in range(3):
+        torch.manual_seed(12 + rank + i)
         # check multiple random inputs as random exact occurencies are rare
         metric_devices = ["cpu"]
         if device.type != "xla":
diff --git a/tests/ignite/metrics/test_metric.py b/tests/ignite/metrics/test_metric.py
@@ -710,7 +710,7 @@ def _test_creating_on_xla_fails(device):
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
-def test_distrib_nccl_gpu(distributed_context_single_node_nccl):
+def test_distrib_nccl_gpu(distributed):
     device = idist.device()
     _test_distrib_sync_all_reduce_decorator(device)
     _test_invalid_sync_all_reduce(device)
diff --git a/tests/ignite/metrics/test_metrics_lambda.py b/tests/ignite/metrics/test_metrics_lambda.py
@@ -8,7 +8,7 @@
 
 import ignite.distributed as idist
 from ignite.engine import Engine
-from ignite.metrics import Metric, MetricsLambda, Precision, Recall
+from ignite.metrics import Accuracy, Metric, MetricsLambda, Precision, Recall
 
 
 class ListGatherMetric(Metric):
@@ -270,6 +270,34 @@ def Fbeta(r, p, beta):
         assert precision_true == approx(precision), f"{precision_true} vs {precision}"
         assert recall_true == approx(recall), f"{recall_true} vs {recall}"
 
+    metric_state = F1.state_dict()
+    F1.reset()
+    F1.load_state_dict(metric_state)
+    f1_value = F1.compute()
+    assert f1_value == state.metrics["f1"]
+
+
+def test_load_state_dict():
+    acc = Accuracy()
+    error = 1.0 - acc
+
+    acc.update(
+        (
+            torch.randint(0, 2, size=(8,)),
+            torch.randint(0, 2, size=(8,)),
+        )
+    )
+
+    e = error.compute()
+    a = acc.compute()
+    assert 1.0 - a == e
+
+    metric_state = error.state_dict()
+    error.reset()
+    error.load_state_dict(metric_state)
+    e2 = error.compute()
+    assert e2 == e
+
 
 def test_state_metrics():
     y_pred = torch.randint(0, 2, size=(15, 10, 4)).float()
diff --git a/tests/ignite/metrics/test_running_average.py b/tests/ignite/metrics/test_running_average.py
@@ -113,6 +113,24 @@ def _(engine):
     assert (torch.tensor(metric_acc_running_averages) == accuracy_running_averages).all()
     assert (torch.tensor(metric_loss_running_averages) == loss_running_averages).all()
 
+    metric_state = acc_metric.state_dict()
+    saved__value = acc_metric._value
+    saved_src__num_correct = acc_metric.src._num_correct
+    saved_src__num_examples = acc_metric.src._num_examples
+    acc_metric.reset()
+    acc_metric.load_state_dict(metric_state)
+    assert acc_metric._value == saved__value
+    assert acc_metric.src._num_examples == saved_src__num_examples
+    assert (acc_metric.src._num_correct == saved_src__num_correct).all()
+
+    metric_state = avg_output.state_dict()
+    saved__value = avg_output._value
+    assert avg_output.src is None
+    avg_output.reset()
+    avg_output.load_state_dict(metric_state)
+    assert avg_output._value == saved__value
+    assert avg_output.src is None
+
 
 def test_integration_epochwise():
     torch.manual_seed(10)
