An improvement and updating an example

examples/mnist/mnist_save_resume_engine.py

@@ -13,7 +13,7 @@
 
 from ignite.engine import create_supervised_evaluator, create_supervised_trainer, Events
 from ignite.handlers import Checkpoint, DiskSaver
-from ignite.metrics import Accuracy, Loss
+from ignite.metrics import Accuracy, Loss, RunningAverage
 from ignite.utils import manual_seed
 
 try:
@@ -162,24 +162,26 @@ def run(
     if deterministic:
         tqdm.write("Setup deterministic trainer")
     trainer = create_supervised_trainer(model, optimizer, criterion, device=device, deterministic=deterministic)
+    running_loss = RunningAverage(output_transform=lambda x: x)
+    running_loss.attach(trainer, "rloss")
 
-    evaluator = create_supervised_evaluator(
-        model, metrics={"accuracy": Accuracy(), "nll": Loss(criterion)}, device=device
-    )
+    metrics = {"accuracy": Accuracy(), "nll": Loss(criterion)}
+    evaluator = create_supervised_evaluator(model, metrics, device)
 
     # Apply learning rate scheduling
     @trainer.on(Events.EPOCH_COMPLETED)
     def lr_step(engine):
         lr_scheduler.step()
 
-    pbar = tqdm(initial=0, leave=False, total=len(train_loader), desc=f"Epoch {0} - loss: {0:.4f} - lr: {lr:.4f}")
+    pbar = tqdm(initial=0, leave=False, total=len(train_loader), desc=f"Epoch {0} - run. loss: {0:.4f} - lr: {lr:.4f}")
 
     @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
     def log_training_loss(engine):
         lr = optimizer.param_groups[0]["lr"]
-        pbar.desc = f"Epoch {engine.state.epoch} - loss: {engine.state.output:.4f} - lr: {lr:.4f}"
+        rloss = engine.state.metrics["rloss"]
+        pbar.desc = f"Epoch {engine.state.epoch} - run. loss: {rloss:.4f} - lr: {lr:.4f}"
         pbar.update(log_interval)
-        writer.add_scalar("training/loss", engine.state.output, engine.state.iteration)
+        writer.add_scalar("training/running_loss", rloss, engine.state.iteration)
         writer.add_scalar("lr", lr, engine.state.iteration)
 
     if crash_iteration > 0:
@@ -222,7 +224,14 @@ def log_validation_results(engine):
         writer.add_scalar("valdation/avg_accuracy", avg_accuracy, engine.state.epoch)
 
     # Setup object to checkpoint
-    objects_to_checkpoint = {"trainer": trainer, "model": model, "optimizer": optimizer, "lr_scheduler": lr_scheduler}
+    objects_to_checkpoint = {
+        "trainer": trainer,
+        "model": model,
+        "optimizer": optimizer,
+        "lr_scheduler": lr_scheduler,
+        "train_running_loss": running_loss,
+        "metrics": metrics,
+    }
     training_checkpoint = Checkpoint(
         to_save=objects_to_checkpoint,
         save_handler=DiskSaver(log_dir, require_empty=False),

ignite/handlers/checkpoint.py

@@ -476,7 +476,10 @@ def _setup_checkpoint(self) -> Dict[str, Dict[Any, Any]]:
                     obj.consolidate_state_dict(to=self.save_on_rank)
                     if self.save_on_rank != idist.get_rank():
                         continue
-                checkpoint[k] = obj.state_dict()
+                if is_dict_of_serializable_objects(obj):
+                    checkpoint[k] = {k: v.state_dict() for k, v in obj.items()}
+                else:
+                    checkpoint[k] = obj.state_dict()
         return checkpoint
 
     @staticmethod
@@ -532,8 +535,8 @@ def setup_filename_pattern(
 
     @staticmethod
     def _check_objects(objs: Mapping, attr: str) -> None:
-        for k, obj in objs.items():
-            if not hasattr(obj, attr):
+        for _, obj in objs.items():
+            if not hasattr(obj, attr) and not is_dict_of_serializable_objects(obj):
                 raise TypeError(f"Object {type(obj)} should have `{attr}` method")
 
     @staticmethod
@@ -611,6 +614,9 @@ def _load_object(obj: Any, chkpt_obj: Any) -> None:
                 obj = obj.module
             if isinstance(obj, torch.nn.Module):
                 obj.load_state_dict(chkpt_obj, strict=is_state_dict_strict)
+            elif is_dict_of_serializable_objects(obj):
+                for k, v in obj.items():
+                    v.load_state_dict(chkpt_obj[k])
             else:
                 obj.load_state_dict(chkpt_obj)
 
@@ -1002,3 +1008,7 @@ def __call__(self, engine: Engine, to_save: Mapping):  # type: ignore
         self._check_objects(to_save, "state_dict")
         self.to_save = to_save
         super(ModelCheckpoint, self).__call__(engine)
+
+
+def is_dict_of_serializable_objects(obj: Any) -> bool:
+    return isinstance(obj, Mapping) and all([hasattr(v, "state_dict") for v in obj.values()])