Repository: pytorch/ignite

Issue:
Does ignite support to load weights for different num_classes?
## ❓ Questions/Help/Support
Hi @vfdev-5 ,

Does ignite Checkpoint handler support the case that pre-trained model has different number output with the model load to?
For example, the weights in model checkpoint is for 1000 classes output, and now user defines the network with 3 classes output and use ignite handler to load pretrained model for the weights.
Thanks.
<!-- If you have a question, please checkout other issues labelled as "question" (https://github.com/pytorch/ignite/issues?q=is%3Aissue+is%3Aopen+label%3Aquestion) -->
<!-- If your question is not listed above or you need a support or help, please provide a description-->
<!-- If you prefer, you can also reach out us on [Discussion Forum](https://discuss.pytorch.org/) adding tag "ignite". -->



Discussion:
Hi @Nic-Ma 
I think this is currently not supported, because we simply load the checkpoint and don't do any checks on the Tensor sizes. 
We should allow kwargs in the `load_objects` function and you would have to pass everything except the last layer as `to_load`. I think implementing a general interface for that is not that easy, but I'll think on that.
Hi @Nic-Ma ,

yes, as @justusschock said this is not supported and probably the reason for that is related to the fact that pretrained weights loading in most of the cases provided by model definition (e.g. in torchvision `resnet18(pretrained=True)`). In that case, user's modification can be easily done a posteriori :
```python
model = resnet18(pretrained=True)
model.fc = nn.Linear(512, 3)
``` 
@vfdev-5 we should nevertheless think about kwargs for the load_objects function, to allow kwargs for both, `torch.load` (for example the `map_location`) and the `load_state_dict` (for example the `strict`).
Yes, I agree!
Hi @vfdev-5 and @justusschock ,

Thanks very much for your sharing. Looking forward to this feature.
It's really useful for us, especially to provide "Transfer Learning" function in our product.
Currently, we can't exclude FC layer when loading weights, users raised this request.
Thanks.

Resolving patch:
diff --git a/ignite/handlers/checkpoint.py b/ignite/handlers/checkpoint.py
--- a/ignite/handlers/checkpoint.py
+++ b/ignite/handlers/checkpoint.py
@@ -242,7 +242,7 @@ def _check_objects(objs: Mapping, attr: str) -> None:
                 raise TypeError("Object {} should have `{}` method".format(type(obj), attr))
 
     @staticmethod
-    def load_objects(to_load: Mapping, checkpoint: Mapping) -> None:
+    def load_objects(to_load: Mapping, checkpoint: Mapping, **kwargs) -> None:
         """Helper method to apply `load_state_dict` on the objects from `to_load` using states from `checkpoint`.
 
         Args:
@@ -250,22 +250,32 @@ def load_objects(to_load: Mapping, checkpoint: Mapping) -> None:
             checkpoint (Mapping): a dictionary with state_dicts to load, e.g. `{"model": model_state_dict,
                 "optimizer": opt_state_dict}`. If `to_load` contains a single key, then checkpoint can contain directly
                 corresponding state_dict.
+            **kwargs: Keyword arguments accepted for `nn.Module.load_state_dict()`. Passing `strict=False` enables
+                the user to load part of the pretrained model (useful for example, in Transfer Learning)
         """
         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(kwargs) > 1 or any(k for k in kwargs.keys() if k not in ["strict"]):
+            warnings.warn("kwargs contains keys other than strict and these will be ignored")
+
+        is_state_dict_strict = kwargs.get("strict", True)
         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)
+                obj.load_state_dict(checkpoint, strict=is_state_dict_strict)
                 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))
-            obj.load_state_dict(checkpoint[k])
+            if isinstance(obj, torch.nn.Module):
+                obj.load_state_dict(checkpoint[k], strict=is_state_dict_strict)
+            else:
+                obj.load_state_dict(checkpoint[k])
 
 
 class DiskSaver:


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
@@ -22,15 +22,25 @@ def forward(self, x):
         return self.net(x)
 
 
+class DummyPretrainedModel(nn.Module):
+    def __init__(self):
+        super(DummyPretrainedModel, self).__init__()
+        self.features = nn.Linear(4, 2, bias=False)
+        self.fc = nn.Linear(2, 1)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.fc(x)
+        return x
+
+
 def test_checkpoint_wrong_input():
 
     with pytest.raises(TypeError, match=r"Argument `to_save` should be a dictionary"):
-        Checkpoint(
-            12, lambda x: x, "prefix",
-        )
+        Checkpoint(12, lambda x: x, "prefix")
 
     with pytest.raises(TypeError, match=r"Argument `to_save` should be a dictionary"):
-        Checkpoint([12,], lambda x: x, "prefix")
+        Checkpoint([12], lambda x: x, "prefix")
 
     with pytest.raises(ValueError, match=r"No objects to checkpoint."):
         Checkpoint({}, lambda x: x, "prefix")
@@ -675,7 +685,7 @@ def update_fn(engine, batch):
     handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
 
     engine.add_event_handler(
-        Events.EPOCH_COMPLETED, handler, {"model": model, "optimizer": optim, "lr_scheduler": lr_scheduler,}
+        Events.EPOCH_COMPLETED, handler, {"model": model, "optimizer": optim, "lr_scheduler": lr_scheduler}
     )
     engine.run([0], max_epochs=4)
 
@@ -742,20 +752,14 @@ def test_checkpoint_load_objects():
 def test_checkpoint_load_objects_from_saved_file(dirname):
     def _get_single_obj_to_save():
         model = DummyModel()
-        to_save = {
-            "model": model,
-        }
+        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,
-        }
+        to_save = {"model": model, "optimizer": optim, "lr_scheduler": lr_scheduler}
         return to_save
 
     trainer = Engine(lambda e, b: None)
@@ -798,6 +802,33 @@ def _get_multiple_objs_to_save():
     Checkpoint.load_objects(to_save, loaded_objects)
 
 
+def test_load_checkpoint_with_different_num_classes(dirname):
+    model = DummyPretrainedModel()
+    to_save_single_object = {"model": model}
+
+    trainer = Engine(lambda e, b: None)
+    trainer.state = State(epoch=0, iteration=0)
+
+    handler = ModelCheckpoint(dirname, _PREFIX, create_dir=False, n_saved=1)
+    handler(trainer, to_save_single_object)
+
+    fname = handler.last_checkpoint
+    loaded_checkpoint = torch.load(fname)
+
+    to_load_single_object = {"pretrained_features": model.features}
+
+    with pytest.raises(RuntimeError):
+        Checkpoint.load_objects(to_load_single_object, loaded_checkpoint)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        Checkpoint.load_objects(to_load_single_object, loaded_checkpoint, strict=False, blah="blah")
+
+    loaded_weights = to_load_single_object["pretrained_features"].state_dict()["weight"]
+
+    assert torch.all(model.state_dict()["features.weight"].eq(loaded_weights))
+
+
 def test_disksaver_wrong_input(dirname):
 
     with pytest.raises(ValueError, match=r"Directory path '\S+' is not found"):