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:

Patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
index 4f806cd5..6e66a2b8 100644
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -236,11 +236,20 @@ class Checkpoint:
         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
index 208698bc..a853df7c 100644
--- 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"):