How to save / load part of checkpoint?

[Nic-Ma]

❓ Questions/Help/Support

Hi @vfdev-5 ,

I have a feature request from Clara users for transfer learning: I trained a model with 4 output classes, then I want to load the checkpoint for a 3 classes task.
How to save weights except for the last output layer in ignite Checkpoint handler? Or skip the last layer when loading checkpoint in the handler?

Thanks in advance.

[Nic-Ma]

I tested with strict=False in this line:
https://github.com/pytorch/ignite/blob/master/ignite/handlers/checkpoint.py#L530
But seems it can't skip "same layer name but different size"?

Here is the test result of "loading the UNet ckpt trained on spleen dataset to prostate dataset":

  File "/workspace/data/medical/MONAI/monai/handlers/checkpoint_loader.py", line 89, in __call__
    Checkpoint.load_objects(to_load=self.load_dict, checkpoint=checkpoint, strict=self.strict)
  File "/opt/conda/lib/python3.6/site-packages/ignite/handlers/checkpoint.py", line 556, in load_objects
    obj.load_state_dict(checkpoint, strict=is_state_dict_strict)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1052, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for UNet:
      size mismatch for model.0.conv.unit0.conv.weight: copying a param with shape torch.Size([16, 1, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([16, 2, 3, 3, 3]).
      size mismatch for model.0.residual.weight: copying a param with shape torch.Size([16, 1, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([16, 2, 3, 3, 3]).
      size mismatch for model.2.0.conv.weight: copying a param with shape torch.Size([32, 2, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([32, 3, 3, 3, 3]).
      size mismatch for model.2.0.conv.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([3]).
      size mismatch for model.2.0.adn.N.weight: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([3]).
      size mismatch for model.2.0.adn.N.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([3]).
      size mismatch for model.2.0.adn.N.running_mean: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([3]).
      size mismatch for model.2.0.adn.N.running_var: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([3]).
      size mismatch for model.2.1.conv.unit0.conv.weight: copying a param with shape torch.Size([2, 2, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([3, 3, 3, 3, 3]).
      size mismatch for model.2.1.conv.unit0.conv.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([3]).

Thanks.

[vfdev-5]

Hi @Nic-Ma, thanks for the question ! Yes, that's true that current version of ignite does not support this option when a layer is present in the saved file and in the model but they have different shapes. Personnally, I think this could be error prone.

I'd like to better understand the steps:

• train a model for example with 4 classes and save the best model only in a .pt file (e.g. best_model_4.pt)
• finetune a similar model with a smaller head, for example, with 3 classes and use best_model_4.pt to load pretrained weights

I think in this case loading should be done in a separate way:

model_3 = get_model_with_3_classes()
state_dict_model_4 = torch.load("best_model_4.pt")

@trainer.on(Events.STARTED, model_3, state_dict_model_4)
def permissive_model_loader(model, state_dict):
    this_state_dict = model.state_dict()
    matched_state_dict = {
        k: v for k, v in state_dict.items()
        if k in this_state_dict and v.shape == this_state_dict[k].shape
    }
    model.load_state_dict(matched_state_dict, strict=False)

trainer.run(...)

Another option could be to define manually what to load and reuse monai's CheckpointLoader :

model_3 = get_model_with_3_classes()
# model_3.features == model_4.features # common part of 2 models

loader = CheckpointLoader(
    "best_model_4.pt", 
    {"network": model_3.features}
)
...

What do you think ?

[Nic-Ma]

Hi @vfdev-5 ,

Thanks for your detailed solutions!
I prefer the first solution because we may not know the exact layers affected by class numbers, for example, maybe a different input channel number will affect several convolution layers.
Maybe we can enhance the MONAI CheckpointLoader with the strick_shape option later, I submitted ticket:
Project-MONAI/MONAI#1939

Thanks.

[vfdev-5]

@Nic-Ma sounds good to me, however, I think we should be careful with such permissive loading. For example, your above case where base UNet has 2 input channels for the first convolutions (input has 2 channels) and finetuning model works on the input with 1 channel. Thus, we load middle layers and keep random top and head layers:

Layers	Trained UNet (2 -> 4)	Finetuning UNet (1 -> 3) after loading
0.conv	Trained	Random
0.adn	Trained	Same as left
1.*	Trained	Same as left
2.*	Trained	Random

In this case, users who tweak lr rate for pretrained layers should really know what they are doing. We can imagine the impact if accidently top layer is random and is frozen by user thinking that it is pretrained...

Btw, do we still "freeze" pretrained layers in this case when top is random ?

[Nic-Ma]

Hi @vfdev-5 ,

Good suggestions.
I think this should be advanced feature and users can set "loading layers" and "finetune LR" separately.
Currently, we have the API to set different LR for different layers in MONAI:
https://github.com/Project-MONAI/MONAI/blob/master/monai/optimizers/utils.py#L21

Thanks.

[vfdev-5]

I think we can close this issue as solved.