Can't load UNETR pretrained model correctly

[chaoscls]

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!

[tangy5]

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.

[RR-N]

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.

[Cocofeat]

I fix it with the monai==0.7.0