Repository: pfnet/pytorch-pfn-extras
Issue: GradScaler state_dict not saved in snapshot

GradScaler objects can be managed by ppe.training.manager._BaseExtensionManager so that state_dict can be saved

Changed files:
pytorch_pfn_extras/training/_manager_protocol.py, pytorch_pfn_extras/training/manager.py, tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py, pytorch_pfn_extras/training/__init__.py, tests/pytorch_pfn_extras_tests/training_tests/test_manager.py, pytorch_pfn_extras/engine.py

diff --git a/pytorch_pfn_extras/engine.py b/pytorch_pfn_extras/engine.py
index 3db3eea..0b6683d 100644
--- a/pytorch_pfn_extras/engine.py
+++ b/pytorch_pfn_extras/engine.py
@@ -14,6 +14,7 @@ from typing import (
 import pytorch_pfn_extras.handler as handler_module
 import torch
 from pytorch_pfn_extras.runtime import runtime_registry
+from pytorch_pfn_extras.training import StateObjectProtocol
 from pytorch_pfn_extras.training._transform_model import default_transform_model
 
 if TYPE_CHECKING:
@@ -106,10 +107,21 @@ 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
+
     # TODO(kmaehashi): deprecate specifying 'runtime' key in options
     runtime_options = dict(
         runtime_options if runtime_options else options.pop("runtime", {})
     )
+
+    for key, value in runtime_options.items():
+        if isinstance(value, StateObjectProtocol):
+            state_objects[f"runtime_options_{key}"] = value
+
     logic = handler_module.Logic() if logic is None else logic
     handler_class = handler_class if handler_class else handler_module.Handler
 
@@ -139,6 +151,7 @@ def create_trainer(
         writer=writer,
         transform_model=transform_model,
         profile=profile,
+        state_objects=state_objects,
         **kwargs,
     )
 
diff --git a/pytorch_pfn_extras/training/__init__.py b/pytorch_pfn_extras/training/__init__.py
index 8fae0df..4289657 100644
--- a/pytorch_pfn_extras/training/__init__.py
+++ b/pytorch_pfn_extras/training/__init__.py
@@ -3,6 +3,7 @@ from pytorch_pfn_extras.training._evaluator import DistributedEvaluator  # NOQA
 from pytorch_pfn_extras.training._evaluator import Evaluator  # NOQA
 from pytorch_pfn_extras.training._manager_protocol import (  # NOQA
     ExtensionsManagerProtocol,
+    StateObjectProtocol,
 )
 from pytorch_pfn_extras.training._trainer import Trainer  # NOQA
 from pytorch_pfn_extras.training.extension import PRIORITY_EDITOR  # NOQA
diff --git a/pytorch_pfn_extras/training/_manager_protocol.py b/pytorch_pfn_extras/training/_manager_protocol.py
index 624acff..08f88d5 100644
--- a/pytorch_pfn_extras/training/_manager_protocol.py
+++ b/pytorch_pfn_extras/training/_manager_protocol.py
@@ -1,4 +1,11 @@
-from typing import TYPE_CHECKING, Mapping, Optional
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Mapping,
+    Optional,
+    runtime_checkable,
+)
 
 import torch
 from typing_extensions import Protocol
@@ -72,3 +79,12 @@ class ExtensionsManagerProtocol(Protocol):
     @property
     def observation(self) -> "reporting.Observation":
         ...
+
+
+@runtime_checkable
+class StateObjectProtocol(Protocol):
+    def state_dict(self) -> Dict[str, Any]:
+        ...
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        ...
diff --git a/pytorch_pfn_extras/training/manager.py b/pytorch_pfn_extras/training/manager.py
index 3ba206e..487c282 100644
--- a/pytorch_pfn_extras/training/manager.py
+++ b/pytorch_pfn_extras/training/manager.py
@@ -1,4 +1,4 @@
-import collections
+import collections.abc
 import contextlib
 import copy
 import time
@@ -18,6 +18,7 @@ import pytorch_pfn_extras
 import torch
 from pytorch_pfn_extras import reporting, writing
 from pytorch_pfn_extras.profiler import record
+from pytorch_pfn_extras.training import StateObjectProtocol
 from pytorch_pfn_extras.training import _util as util_module
 from pytorch_pfn_extras.training import extension as extension_module
 from pytorch_pfn_extras.training import trigger as trigger_module
@@ -99,6 +100,9 @@ class _ManagerProxy:
         return self._manager.observation
 
 
+_default_state_objects: Dict[str, StateObjectProtocol] = {}
+
+
 class _BaseExtensionsManager:
     """
     Keeps track of the extensions and the current status
@@ -117,6 +121,7 @@ class _BaseExtensionsManager:
         stop_trigger: "trigger_module.TriggerLike" = None,
         transform_model: _TransformModel = default_transform_model,
         enable_profile: bool = False,
+        state_objects: Dict[str, StateObjectProtocol] = _default_state_objects,
     ) -> None:
         if extensions is None:
             extensions = []
@@ -178,6 +183,7 @@ class _BaseExtensionsManager:
             self.extend(ext)
 
         self._enable_profile = enable_profile
+        self._state_objects = state_objects
         # Initialize the writer
         self.writer.initialize(self.out)
 
@@ -493,6 +499,10 @@ class _BaseExtensionsManager:
             name: self._optimizers[name].state_dict()
             for name in self._optimizers
         }
+        to_save["state_objects"] = {
+            name: self._state_objects[name].state_dict()
+            for name in self._state_objects
+        }
         to_save["extensions"] = {
             name: self._extensions[name].state_dict()
             for name in self._extensions
@@ -538,6 +548,12 @@ class _BaseExtensionsManager:
         for name in self._optimizers:
             self._optimizers[name].load_state_dict(to_load["optimizers"][name])
 
+        if "state_objects" in to_load:
+            for name in self._state_objects:
+                self._state_objects[name].load_state_dict(  # type: ignore[no-untyped-call]
+                    to_load["state_objects"][name]
+                )
+
         for name in self._extensions:
             self._extensions[name].load_state_dict(to_load["extensions"][name])
 
@@ -579,6 +595,7 @@ class ExtensionsManager(_BaseExtensionsManager):
         writer: Optional[writing.Writer] = None,
         transform_model: _TransformModel = lambda n, x: x,
         enable_profile: bool = False,
+        state_objects: Dict[str, StateObjectProtocol] = _default_state_objects,
     ) -> None:
         super().__init__(
             models,
@@ -590,6 +607,7 @@ class ExtensionsManager(_BaseExtensionsManager):
             stop_trigger,
             transform_model,
             enable_profile,
+            state_objects,
         )
         if iters_per_epoch < 1:
             raise ValueError(
@@ -683,6 +701,7 @@ class IgniteExtensionsManager(_BaseExtensionsManager):
         out_dir: str = "result",
         writer: Optional[writing.Writer] = None,
         enable_profile: bool = False,
+        state_objects: Dict[str, StateObjectProtocol] = _default_state_objects,
     ) -> None:
         import ignite
 
@@ -703,6 +722,7 @@ class IgniteExtensionsManager(_BaseExtensionsManager):
             out_dir,
             writer,
             enable_profile=enable_profile,
+            state_objects=state_objects,
         )
         self.engine = engine
         self._start_epoch = 0  # Used to correctly restore snapshots
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py b/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
index f0aa1b9..da469f5 100644
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_manager.py
@@ -158,6 +158,7 @@ def test_extensions_manager_state_dict():
     model_state_dict = object()
     optimizer_state_dict = object()
     extension_state_dict = object()
+    misc_state_dict = object()
     max_epochs = 5
     iters_per_epoch = 4
     passed_iteration = 11
@@ -167,6 +168,7 @@ def test_extensions_manager_state_dict():
         {"optimizer_name": _StateDictObj(state_dict=optimizer_state_dict)},
         max_epochs,
         iters_per_epoch=iters_per_epoch,
+        state_objects={"misc": _StateDictObj(state_dict=misc_state_dict)},
     )
 
     manager.extend(
@@ -186,6 +188,7 @@ def test_extensions_manager_state_dict():
         "_start_iteration": passed_iteration,
         "models": {"model_name": model_state_dict},
         "optimizers": {"optimizer_name": optimizer_state_dict},
+        "state_objects": {"misc": misc_state_dict},
         "extensions": {
             "extension_name": {
                 "extension": extension_state_dict,
@@ -354,6 +357,7 @@ def test_ignite_extensions_manager_state_dict():
         "_epoch_length": iters_per_epoch,
         "models": {"model_name": model_state_dict},
         "optimizers": {"optimizer_name": optimizer_state_dict},
+        "state_objects": {},
         "extensions": {
             "extension_name": {
                 "extension": extension_state_dict,
@@ -404,6 +408,7 @@ def test_extensions_manager_with_plain_model_and_optimizer():
         "_start_iteration": 0,
         "models": {"main": model_state_dict},
         "optimizers": {"main": optimizer_state_dict},
+        "state_objects": {},
         "extensions": {},
     }
 
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py b/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
index 916ce50..cfcf482 100644
--- a/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/test_trainer.py
@@ -6,6 +6,7 @@ import pytest
 import pytorch_pfn_extras as ppe
 import torch
 from pytorch_pfn_extras import engine, training
+from pytorch_pfn_extras.training import triggers
 from torch import nn
 from torch.nn import functional as F
 
@@ -359,7 +360,13 @@ def test_train_result_equal(device, path):
         assert torch.equal(a, e)
 
 
-def _compare_states(s1, s2):
+def _compare_states(s1, s2, strict=False):
+    def allclose(a, b):
+        if strict:
+            return (a == b).all()
+        else:
+            return torch.allclose(a, b)
+
     if isinstance(s1, dict):
         keys = s1.keys()
         if set(keys) != set(s2.keys()):
@@ -380,7 +387,7 @@ def _compare_states(s1, s2):
                 return False
             all_equal = all_equal and _compare_states(s1[k], s2[k])
         elif isinstance(s1[k], torch.Tensor):
-            all_equal = all_equal and torch.allclose(s1[k], s2[k])
+            all_equal = all_equal and allclose(s1[k], s2[k])
         else:
             all_equal = all_equal and s1[k] == s2[k]
         if not all_equal:
@@ -389,19 +396,27 @@ def _compare_states(s1, s2):
 
 
 class TestTrainerState:
-    def _get_trainer(self, epochs, out_dir):
+    def _get_trainer(
+        self,
+        epochs,
+        out_dir,
+        extensions=None,
+        options=None,
+        device="cpu",
+    ):
         model = MyModel()
-        ppe.to(model, "cpu")
+        ppe.to(model, device)
         model_with_loss = MyModelWithLossFn(model)
         optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
-        extensions = _make_extensions()
+        extensions = extensions or _make_extensions()
         trainer = engine.create_trainer(
             model_with_loss,
             optimizer,
-            20,
-            device="cpu",
+            epochs,
+            device=device,
             extensions=extensions,
             out_dir=out_dir,
+            options=options,
         )
         return trainer
 
@@ -458,6 +473,149 @@ class TestTrainerState:
         assert new_trainer.epoch == 20
         assert _compare_states(trainer.state_dict(), new_trainer.state_dict())
 
+    def test_trainer_autoload_training_results_consistency(self, path):
+        snapshot_epoch = 10
+        training_epoch = 20
+        trainer = self._get_trainer(training_epoch, path)
+        data = torch.utils.data.DataLoader(
+            [
+                (
+                    torch.rand(
+                        20,
+                    ),
+                    torch.rand(
+                        10,
+                    ),
+                )
+                for i in range(10)
+            ]
+        )
+        trainer.extend(
+            ppe.training.extensions.snapshot(),
+            trigger=triggers.ManualScheduleTrigger([snapshot_epoch], "epoch"),
+        )
+        trainer.run(data)
+        new_trainer = self._get_trainer(training_epoch, path)
+        new_trainer.extend(ppe.training.extensions.snapshot(autoload=True))
+        new_trainer._setup_manager(len(data))
+        assert new_trainer.epoch == snapshot_epoch
+        new_trainer.run(data)
+        assert new_trainer.epoch == training_epoch
+        print(trainer.state_dict().keys())
+        trainer_state_dict = trainer.state_dict()
+        new_trainer_state_dict = new_trainer.state_dict()
+        assert _compare_states(
+            trainer_state_dict["models"],
+            new_trainer_state_dict["models"],
+            strict=True,
+        )
+
+    @pytest.mark.gpu
+    def test_trainer_autoload_training_results_consistency_with_gpu(self, path):
+        if not torch.cuda.is_available():
+            pytest.skip()
+        snapshot_epoch = 10
+        training_epoch = 20
+        trainer = self._get_trainer(training_epoch, path, device="cuda")
+        data = torch.utils.data.DataLoader(
+            [
+                (
+                    torch.rand(
+                        20,
+                    ),
+                    torch.rand(
+                        10,
+                    ),
+                )
+                for i in range(10)
+            ]
+        )
+        trainer.extend(
+            ppe.training.extensions.snapshot(),
+            trigger=triggers.ManualScheduleTrigger([snapshot_epoch], "epoch"),
+        )
+        trainer.run(data)
+        new_trainer = self._get_trainer(training_epoch, path, device="cuda")
+        new_trainer.extend(ppe.training.extensions.snapshot(autoload=True))
+        new_trainer._setup_manager(len(data))
+        assert new_trainer.epoch == snapshot_epoch
+        new_trainer.run(data)
+        assert new_trainer.epoch == training_epoch
+        print(trainer.state_dict().keys())
+        trainer_state_dict = trainer.state_dict()
+        new_trainer_state_dict = new_trainer.state_dict()
+        assert _compare_states(
+            trainer_state_dict["models"],
+            new_trainer_state_dict["models"],
+            strict=True,
+        )
+
+    @pytest.mark.gpu
+    def test_trainer_autoload_training_results_consistency_with_gradscaler(
+        self, path
+    ):
+        if not torch.cuda.is_available():
+            pytest.skip()
+        snapshot_epoch = 10
+        training_epoch = 20
+        grad_scaler = torch.cuda.amp.grad_scaler.GradScaler(
+            init_scale=2**48, growth_interval=2
+        )
+        trainer = self._get_trainer(
+            training_epoch,
+            path,
+            options={
+                "grad_scaler": grad_scaler,
+            },
+            device="cuda",
+        )
+        data = torch.utils.data.DataLoader(
+            [
+                (
+                    torch.rand(
+                        20,
+                    ),
+                    torch.rand(
+                        10,
+                    ),
+                )
+                for i in range(10)
+            ]
+        )
+        trainer.extend(
+            ppe.training.extensions.snapshot(),
+            trigger=triggers.ManualScheduleTrigger([snapshot_epoch], "epoch"),
+        )
+
+        trainer.run(data)
+
+        new_grad_scaler = torch.cuda.amp.grad_scaler.GradScaler(
+            init_scale=2**48, growth_interval=2
+        )
+        new_trainer = self._get_trainer(
+            training_epoch,
+            path,
+            options={
+                "grad_scaler": new_grad_scaler,
+            },
+            device="cuda",
+        )
+        new_trainer.extend(ppe.training.extensions.snapshot(autoload=True))
+        new_trainer._setup_manager(len(data))
+        assert new_trainer.epoch == snapshot_epoch
+
+        new_trainer.run(data)
+        assert new_trainer.epoch == training_epoch
+
+        trainer_state_dict = trainer.state_dict()
+        new_trainer_state_dict = new_trainer.state_dict()
+
+        assert _compare_states(
+            trainer_state_dict["models"],
+            new_trainer_state_dict["models"],
+            strict=True,
+        )
+
 
 class MyModelWithLossDictOutput(torch.nn.Module):
     def __init__(self, model):