Repository: pfnet/pytorch-pfn-extras

Issue:
Wrong counters when restoring snapshots
When loading a snapshot of a finished training using `autoload=True` The training runs an additional epoch while it should not run at all.


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/training/manager.py b/pytorch_pfn_extras/training/manager.py
--- a/pytorch_pfn_extras/training/manager.py
+++ b/pytorch_pfn_extras/training/manager.py
@@ -76,6 +76,7 @@ def __init__(
         self._out = out_dir
         self.writer = writer
         self.reporter = Reporter()
+        self._start_extensions_called = False
 
         if not isinstance(models, dict):
             if not isinstance(models, torch.nn.Module):
@@ -106,12 +107,23 @@ def __init__(
         # Initialize the writer
         self.writer.initialize(self.out)
 
+    @property
+    def iteration(self):
+        self.start_extensions()
+        return self._iteration
+
+    @iteration.setter
+    def iteration(self, value):
+        self._iteration = value
+
     @property
     def models(self):
+        self.start_extensions()
         return self._models
 
     @property
     def optimizers(self):
+        self.start_extensions()
         return self._optimizers
 
     @property
@@ -132,6 +144,7 @@ def epoch_detail(self):
 
     @property
     def stop_trigger(self):
+        self.start_extensions()
         # Trigger is stateful, we close the extensions the first time
         # it evaluates to True, as it won't do it again
         return self._stop_trigger(self)
@@ -160,6 +173,12 @@ def _prepare_for_training(self, start_iteration, iters_per_epoch):
         self._iters_per_epoch = iters_per_epoch
 
     def start_extensions(self):
+        if self._start_extensions_called:
+            # This method must not be called twice or more.
+            return
+        else:
+            self._start_extensions_called = True
+
         exts = self._extensions
         extension_order = sorted(
             exts.keys(),


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_snapshot.py b/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_snapshot.py
--- a/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_snapshot.py
+++ b/tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_snapshot.py
@@ -16,12 +16,11 @@
 from pytorch_pfn_extras import writing
 
 
-def get_trainer(*, out_dir, state_to_load=None):
+def get_trainer(*, out_dir, state_to_load=None, epochs=10):
     model_state_dict = {}
     optimizer_state_dict = {}
     models = {'main': _StateDictModel(state_dict=model_state_dict)}
     optimizers = {'main': _StateDictObj(state_dict=optimizer_state_dict)}
-    epochs = 10  # FIXME
     return training.ExtensionsManager(
         models, optimizers, epochs,
         iters_per_epoch=10,
@@ -361,3 +360,28 @@ def test_model_transformations(path):
     snapshot(manager)
 
     assert model.accessed
+
+
+def test_snapshot_autoload_twice(path):
+    max_epochs = 10
+    iters_per_epoch = 10
+    fmt = 'snapshot_iter_{.iteration}'
+
+    def get_epoch_indices():
+        manager = get_trainer(out_dir=path, epochs=max_epochs)
+        snapshot = extensions.snapshot(filename=fmt, autoload=True)
+        manager.extend(snapshot)
+
+        epoch_indices = []
+        while not manager.stop_trigger:
+            epoch_indices.append(manager.epoch)
+            for it in range(iters_per_epoch):
+                with manager.run_iteration():
+                    time.sleep(0.01)
+        return epoch_indices
+
+    epoch_indices = get_epoch_indices()
+    assert len(epoch_indices) == 10
+
+    epoch_indices = get_epoch_indices()
+    assert len(epoch_indices) == 0