Project-MONAI
Can't load UNETR pretrained model correctly
when I load UNETR pretrained model from https://drive.google.com/file/d/1kR5QuRAuooYcTNLMnMj80Z9IgSs8jtLO/view?usp=sharing, I got error:
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format( RuntimeError: Error(s) in loading state_dict for UNETR: Unexpected key(s) in state_dict: "vit.patch_embedding.cls_token", "encoder2.blocks.0.1.conv3.conv.weight", "encoder2.blocks.1.1.conv3.conv.weight", "encoder3.blocks.0.1.conv3.conv.weight".
How can I fix it? Thanks!
Hi @chaoscls , thanks for the question. The released model should be the one for training BTCV, 13 organs segmentation. I guess if the number of input channels are different (e.g., load the pre-trained model to train other tasks which are not BTCV dataset), there will an issue, especially for the patch embedding weights.
For a quick solution, you can either use strict=false when loading pre-trained weights, or load the model without patch embedding layer keys if the model is used for training other datasets.
For a longer time maintenance, we are working on releasing a model for flexible input/output channels. So that it can be used for other tasks. Note that the current pre-trained model is supervised trained on multi-organ segmentation.
Thank you.
Just noting +1 on the need for more flexible i/o channels -- we were considering UNETR for use on microscopy images and wanted to run a quick test before we commit to training and couldn't due to this same issue.
I fix it with the monai==0.7.0
@Cocofeat this doesn't work for me, could you share your environment,
