Repository: Project-MONAI/MONAI

Issue:
Continue training past checkpoint's max epochs
**Is your feature request related to a problem? Please describe.**
Currently, there is no clean way, to continue training past its initial epochs due to the fact that the state of the trained is being loaded by CheckpointLoader.

**Describe the solution you'd like**
A quick solution would be to assume that if this case appears the code of the user will account for those things (aka just gamma decay, etc) and the CheckpointLoader should only modify the max_epochs after loading the checkpoint data and issue a Warning through a logger.

**Describe alternatives you've considered**
None

**Additional context**
None



Discussion:

Resolving patch:
diff --git a/monai/handlers/checkpoint_loader.py b/monai/handlers/checkpoint_loader.py
--- a/monai/handlers/checkpoint_loader.py
+++ b/monai/handlers/checkpoint_loader.py
@@ -80,5 +80,16 @@ def __call__(self, engine: Engine) -> None:
         """
         checkpoint = torch.load(self.load_path, map_location=self.map_location)
 
+        # save current max epochs setting in the engine, don't overwrite it if larger than max_epochs in checkpoint
+        prior_max_epochs = engine.state.max_epochs
         Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint)
+        if engine.state.epoch > prior_max_epochs:
+            raise ValueError(
+                f"Epoch count ({engine.state.epoch}) in checkpoint is larger than "
+                f"the `engine.state.max_epochs` ({prior_max_epochs}) of engine. To further train from checkpoint, "
+                "construct trainer with `max_epochs` larger than checkpoint's epoch count. "
+                "To use checkpoint for inference, no need to load state_dict for the engine."
+            )
+        engine.state.max_epochs = prior_max_epochs
+
         self.logger.info(f"Restored all variables from {self.load_path}")


Test patch:
diff --git a/tests/test_handler_checkpoint_loader.py b/tests/test_handler_checkpoint_loader.py
--- a/tests/test_handler_checkpoint_loader.py
+++ b/tests/test_handler_checkpoint_loader.py
@@ -16,7 +16,7 @@
 
 import torch
 import torch.optim as optim
-from ignite.engine import Engine
+from ignite.engine import Engine, Events
 
 from monai.handlers import CheckpointLoader, CheckpointSaver
 
@@ -33,15 +33,30 @@ def test_one_save_one_load(self):
         data2["weight"] = torch.tensor([0.2])
         net2.load_state_dict(data2)
         with tempfile.TemporaryDirectory() as tempdir:
-            engine = Engine(lambda e, b: None)
-            CheckpointSaver(save_dir=tempdir, save_dict={"net": net1}, save_final=True).attach(engine)
-            engine.run([0] * 8, max_epochs=5)
-            path = tempdir + "/net_final_iteration=40.pt"
-            engine = Engine(lambda e, b: None)
-            CheckpointLoader(load_path=path, load_dict={"net": net2}).attach(engine)
-            engine.run([0] * 8, max_epochs=1)
+            engine1 = Engine(lambda e, b: None)
+            CheckpointSaver(save_dir=tempdir, save_dict={"net": net1, "eng": engine1}, save_final=True).attach(engine1)
+            engine1.run([0] * 8, max_epochs=5)
+            path = tempdir + "/checkpoint_final_iteration=40.pt"
+            engine2 = Engine(lambda e, b: None)
+            CheckpointLoader(load_path=path, load_dict={"net": net2, "eng": engine2}).attach(engine2)
+
+            @engine2.on(Events.STARTED)
+            def check_epoch(engine: Engine):
+                self.assertEqual(engine.state.epoch, 5)
+
+            engine2.run([0] * 8, max_epochs=8)
             torch.testing.assert_allclose(net2.state_dict()["weight"], torch.tensor([0.1]))
 
+            # test bad case with max_epochs smaller than current epoch
+            engine3 = Engine(lambda e, b: None)
+            CheckpointLoader(load_path=path, load_dict={"net": net2, "eng": engine3}).attach(engine3)
+
+            try:
+                engine3.run([0] * 8, max_epochs=3)
+            except ValueError:
+                self.assertEqual(engine3.state.epoch, 5)
+                self.assertEqual(engine3.state.max_epochs, 5)
+
     def test_two_save_one_load(self):
         logging.basicConfig(stream=sys.stdout, level=logging.INFO)
         net1 = torch.nn.PReLU()