UNet load state_dict error

[YingdiZhang]

Thanks for the great job!
When I train a 3D segmentation model using UNet, I saved the best valiodation model. After that when I try to load the trained state_dict, an error occours and I do not know why:

RuntimeError: Error(s) in loading state_dict for UNet:
Missing key(s) in state_dict: "model.0.conv.unit0.conv.weight", "model.0.conv.unit0.conv.bias", "model.0.conv.unit1.conv.weight", "model.0.conv.unit1.conv.bias", "model.0.residual.weight", "model.0.residual.bias", "model.1.submodule.0.conv.unit0.conv.weight", "model.1.submodule.0.conv.unit0.conv.bias", "model.1.submodule.0.conv.unit1.conv.weight", "model.1.submodule.0.conv.unit1.conv.bias", "model.1.submodule.0.residual.weight", "model.1.submodule.0.residual.bias", "model.1.submodule.1.submodule.0.conv.unit0.conv.weight", "model.1.submodule.1.submodule.0.conv.unit0.conv.bias", "model.1.submodule.1.submodule.0.conv.unit1.conv.weight", "model.1.submodule.1.submodule.0.conv.unit1.conv.bias", "model.1.submodule.1.submodule.0.residual.weight", "model.1.submodule.1.submodule.0.residual.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.bias", "model.1.submodule.1.submodule.1.submodule.0.residual.weight", "model.1.submodule.1.submodule.1.submodule.0.residual.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.residual.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.residual.bias", "model.1.submodule.1.submodule.1.submodule.2.0.conv.weight", "model.1.submodule.1.submodule.1.submodule.2.0.conv.bias", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.weight", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.bias", "model.1.submodule.1.submodule.2.0.conv.weight", "model.1.submodule.1.submodule.2.0.conv.bias", "model.1.submodule.1.submodule.2.1.conv.unit0.conv.weight", "model.1.submodule.1.submodule.2.1.conv.unit0.conv.bias", "model.1.submodule.2.0.conv.weight", "model.1.submodule.2.0.conv.bias", "model.1.submodule.2.1.conv.unit0.conv.weight", "model.1.submodule.2.1.conv.unit0.conv.bias", "model.2.0.conv.weight", "model.2.0.conv.bias", "model.2.1.conv.unit0.conv.weight", "model.2.1.conv.unit0.conv.bias".
Unexpected key(s) in state_dict: "module.model.0.conv.unit0.conv.weight", "module.model.0.conv.unit0.conv.bias", "module.model.0.conv.unit1.conv.weight", "module.model.0.conv.unit1.conv.bias", "module.model.0.residual.weight", "module.model.0.residual.bias", "module.model.1.submodule.0.conv.unit0.conv.weight", "module.model.1.submodule.0.conv.unit0.conv.bias", "module.model.1.submodule.0.conv.unit1.conv.weight", "module.model.1.submodule.0.conv.unit1.conv.bias", "module.model.1.submodule.0.residual.weight", "module.model.1.submodule.0.residual.bias", "module.model.1.submodule.1.submodule.0.conv.unit0.conv.weight", "module.model.1.submodule.1.submodule.0.conv.unit0.conv.bias", "module.model.1.submodule.1.submodule.0.conv.unit1.conv.weight", "module.model.1.submodule.1.submodule.0.conv.unit1.conv.bias", "module.model.1.submodule.1.submodule.0.residual.weight", "module.model.1.submodule.1.submodule.0.residual.bias", "module.model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.weight", "module.model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.bias", "module.model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.weight", "module.model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.bias", "module.model.1.submodule.1.submodule.1.submodule.0.residual.weight", "module.model.1.submodule.1.submodule.1.submodule.0.residual.bias", "module.model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.weight", "module.model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.bias", "module.model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.weight", "module.model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.bias", "module.model.1.submodule.1.submodule.1.submodule.1.submodule.residual.weight", "module.model.1.submodule.1.submodule.1.submodule.1.submodule.residual.bias", "module.model.1.submodule.1.submodule.1.submodule.2.0.conv.weight", "module.model.1.submodule.1.submodule.1.submodule.2.0.conv.bias", "module.model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.weight", "module.model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.bias", "module.model.1.submodule.1.submodule.2.0.conv.weight", "module.model.1.submodule.1.submodule.2.0.conv.bias", "module.model.1.submodule.1.submodule.2.1.conv.unit0.conv.weight", "module.model.1.submodule.1.submodule.2.1.conv.unit0.conv.bias", "module.model.1.submodule.2.0.conv.weight", "module.model.1.submodule.2.0.conv.bias", "module.model.1.submodule.2.1.conv.unit0.conv.weight", "module.model.1.submodule.2.1.conv.unit0.conv.bias", "module.model.2.0.conv.weight", "module.model.2.0.conv.bias", "module.model.2.1.conv.unit0.conv.weight", "module.model.2.1.conv.unit0.conv.bias".

It seems the UNet needs keys like "model.*" but load "module.model."

[Nic-Ma]

Hi @yingdiLee ,

Thanks for your interest here.
I feel that because you used multiple GPUs for training and saved the model, it's common PyTorch question.
You can try to save model.module instead, please refer to:
https://github.com/pytorch/ignite/blob/master/ignite/handlers/checkpoint.py#L581

Thanks.

[Nic-Ma]

Hi @yingdiLee ,

I created a new utility function to help easily save models: #3780

Thanks.