Repository: pytorch/ignite

Issue:
why change checkpoint dict to value if only 1 item?
## 🐛 Bug description
We found there is a logic in checkpoint handler that change the dict to value if only has 1 item.
This logic makes it's hard to restore model with unified load API(_using 'to_load' dict_).
Why did ignite design this logic?

## Environment

 - PyTorch Version (e.g., 1.4): 1.4
 - Ignite Version (e.g., 0.3.0): 0.3.0
 - OS (e.g., Linux): ubuntu
 - How you installed Ignite (`conda`, `pip`, source): pip
 - Python version: 3.6
 - Any other relevant information:
related source code in ignite checkpoint handler line 205:
```py
            if len(checkpoint) == 1:
                for k in checkpoint:
                    name = k
                checkpoint = checkpoint[name]
```

Thanks.




Discussion:
@Nic-Ma the logic is the following:
if user stores a single object in the dictionary `to_save={'model': model}`, then stored filename is `model_XYZ.pth`, where `model` is the key of the single object in `to_save`. If `to_save` contains more objects `{'model': model, 'optimizer': optimizer, ...}`, then filename is `checkpoint_XYZ.pth`. 

Why it is like that ? In most of the cases, if user stores `model`, `optimizer`, `lr_scheduler`, `engine` etc, than this is, IMO, to checkpoint the training. If user stores a single `model` than mostly it is to save the best model.

Hope to answer your question. 

> This logic makes it's hard to restore model with unified load API(using 'to_load' dict).

Could you please detail this point ?

Hi @vfdev-5 ,

Sorry maybe I didn't make the issue clear enough.
The problem is not the filename, it's about this line in CheckpointHandler:
```py
checkpoint = checkpoint[name]
```
It change the **dict** object to **value** if only 1 item existing.
And if use **to_save={'network': net}** to save and **use to_load={'network': net}** to load, it will fail.
Because it didn't save a dict with key "network" if only has this 1 item.
But if use **to_save={'network': net, 'optimizer', optim}** to save,
and use **to_load={'network': net}** to load, it works.
Thanks.
@Nic-Ma thanks for the details. I understand better the problem. IMO, when saving a single object as dictionary seems a bit superficial. But I see your problem with loading method. Let me see if we can improve loading method to cover this case.

To make sure, are you using `Checkpoint.load_objects` method to setup objects?
@vfdev-5 , Yes, we use your loader.

Resolving patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -236,11 +236,20 @@ def load_objects(to_load: Mapping, checkpoint: Mapping) -> None:
         Args:
             to_load (Mapping): a dictionary with objects, e.g. `{"model": model, "optimizer": optimizer, ...}`
             checkpoint (Mapping): a dictionary with state_dicts to load, e.g. `{"model": model_state_dict,
-                "optimizer": opt_state_dict}`
+                "optimizer": opt_state_dict}`. If `to_load` contains a single key, then checkpoint can contain directly
+                corresponding state_dict.
         """
         Checkpoint._check_objects(to_load, "load_state_dict")
         if not isinstance(checkpoint, collections.Mapping):
             raise TypeError("Argument checkpoint should be a dictionary, but given {}".format(type(checkpoint)))
+        if len(to_load) == 1:
+            # single object and checkpoint is directly a state_dict
+            key, obj = list(to_load.items())[0]
+            if key not in checkpoint:
+                obj.load_state_dict(checkpoint)
+                return
+
+        # multiple objects to load
         for k, obj in to_load.items():
             if k not in checkpoint:
                 raise ValueError("Object labeled by '{}' from `to_load` is not found in the checkpoint".format(k))


Test patch:
diff --git a/tests/ignite/handlers/test_checkpoint.py b/tests/ignite/handlers/test_checkpoint.py
--- a/tests/ignite/handlers/test_checkpoint.py
+++ b/tests/ignite/handlers/test_checkpoint.py
@@ -700,7 +700,7 @@ def test_checkpoint_load_objects():
         Checkpoint.load_objects({"a": None}, {"a": None})
 
     model = DummyModel()
-    to_load = {'model': model}
+    to_load = {'model': model, 'another_model': model}
 
     with pytest.raises(ValueError, match=r"from `to_load` is not found in the checkpoint"):
         Checkpoint.load_objects(to_load, {})
@@ -714,6 +714,52 @@ def test_checkpoint_load_objects():
     assert model.state_dict() == model2.state_dict()
 
 
+def test_checkpoint_load_objects_from_saved_file(dirname):
+
+    def _get_single_obj_to_save():
+        model = DummyModel()
+        to_save = {
+            "model": model,
+        }
+        return to_save
+
+    def _get_multiple_objs_to_save():
+        model = DummyModel()
+        optim = torch.optim.SGD(model.parameters(), lr=0.001)
+        lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optim, gamma=0.5)
+        to_save = {
+            "model": model,
+            "optimizer": optim,
+            "lr_scheduler": lr_scheduler,
+        }
+        return to_save
+
+    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
+    trainer = Engine(lambda e, b: None)
+    trainer.state = State(epoch=0, iteration=0)
+    # case: multiple objects
+    to_save = _get_multiple_objs_to_save()
+    handler(trainer, to_save)
+    fname = handler.last_checkpoint
+    assert isinstance(fname, str)
+    assert os.path.join(dirname, _PREFIX) in fname
+    assert os.path.exists(fname)
+    loaded_objects = torch.load(fname)
+    Checkpoint.load_objects(to_save, loaded_objects)
+    os.remove(fname)
+
+    # case: single object
+    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
+    to_save = _get_single_obj_to_save()
+    handler(trainer, to_save)
+    fname = handler.last_checkpoint
+    assert isinstance(fname, str)
+    assert os.path.join(dirname, _PREFIX) in fname
+    assert os.path.exists(fname)
+    loaded_objects = torch.load(fname)
+    Checkpoint.load_objects(to_save, loaded_objects)
+
+
 def test_disksaver_wrong_input(dirname):
 
     with pytest.raises(ValueError, match=r"Directory path '\S+' is not found"):