Merge pull request #716 from linshokaku/dict-state-object

Enable state_dict management for options in the tree structure

Author: emcastillo

pytorch_pfn_extras/engine.py

@@ -3,12 +3,15 @@
     Any,
     Callable,
     Dict,
+    List,
     Mapping,
+    NamedTuple,
     Optional,
     Sequence,
     Tuple,
     Type,
     Union,
+    cast,
 )
 
 import pytorch_pfn_extras.handler as handler_module
@@ -27,6 +30,41 @@
     from pytorch_pfn_extras.training.trigger import TriggerLike
 
 
+def filter_state_objects(
+    args: Any, key_name: str = ""
+) -> List[Tuple[str, StateObjectProtocol]]:
+    if isinstance(args, tuple) and hasattr(args, "_fields"):
+        # namedtuple
+        return filter_state_objects_dict(
+            cast(NamedTuple, args)._asdict(), key_name=key_name
+        )
+    if isinstance(args, dict):
+        return filter_state_objects_dict(args, key_name=key_name)
+    if isinstance(args, (list, tuple)):
+        return sum(
+            [
+                filter_state_objects(v, key_name=f"{key_name}.__{i}__")
+                for i, v in enumerate(args)
+            ],
+            [],
+        )
+    if isinstance(args, StateObjectProtocol):
+        return [(key_name, args)]
+    return []
+
+
+def filter_state_objects_dict(
+    args: Dict[str, Any], key_name: str = "option"
+) -> List[Tuple[str, StateObjectProtocol]]:
+    return sum(
+        [
+            filter_state_objects(v, key_name=f"{key_name}__::__{k}")
+            for k, v in sorted(args.items())
+        ],
+        [],
+    )
+
+
 def create_trainer(
     models: Union[torch.nn.Module, Mapping[str, torch.nn.Module]],
     optimizers: Union[
@@ -109,9 +147,13 @@ def create_trainer(
     options = options.copy() if options else {}
 
     state_objects: Dict[str, StateObjectProtocol] = {}
-    for key, value in options.items():
-        if isinstance(value, StateObjectProtocol):
-            state_objects[f"options_{key}"] = value
+    state_objects_list = filter_state_objects(options)
+    for key, value in state_objects_list:
+        state_objects[f"options_{key}"] = value
+
+    assert len(state_objects_list) == len(
+        state_objects
+    ), "There was a duplicate key name in the flattened options dictionary object."
 
     # TODO(kmaehashi): deprecate specifying 'runtime' key in options
     runtime_options = dict(

tests/pytorch_pfn_extras_tests/training_tests/test_engine.py

@@ -1,3 +1,5 @@
+from typing import Any, Dict
+
 import pytest
 import pytorch_pfn_extras as ppe
 
@@ -75,3 +77,35 @@ def test_extend_after_init(self):
         extension = ppe.training.extensions.LogReport()
         with pytest.raises(RuntimeError, match="cannot extend after"):
             engine.extend(extension)
+
+
+class DummyStateObjects:
+    def __init__(self) -> None:
+        pass
+
+    def state_dict(self) -> Dict[str, Any]:
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        return
+
+
+@pytest.mark.parametrize(
+    "args",
+    [
+        {"a.b": DummyStateObjects(), "a": {"b": DummyStateObjects()}},
+        {"a": DummyStateObjects(), "b": DummyStateObjects()},
+        {
+            "a__.__b": DummyStateObjects(),
+            "__a__": {"__b__": DummyStateObjects()},
+        },
+        {"a::b": DummyStateObjects(), "a": {"b": DummyStateObjects()}},
+        {"a:b": DummyStateObjects(), "a": {"b": DummyStateObjects()}},
+    ],
+)
+def test_filter_state_objects(args) -> None:
+    out = ppe.engine.filter_state_objects(args)
+    key_set = set()
+    for key, _ in out:
+        assert key not in key_set
+        key_set.add(key)