Tiny Model Weight Transfer

[MasalaKimchi]

So I am trying to use the pre-trained weight of the Tiny model (which originally has 14 classes).
But my task has only 5 class predictions, and the code belows gives an error

model = SwinUNETR( img_size=(96, 96, 96), in_channels=1, out_channels=5, feature_size=12, use_checkpoint=True, ).to(device)
model_dict = torch.load("./tiny_model_swinvit.pt")
state_dict = model_dict["state_dict"]
model.load_state_dict(state_dict,`` strict=False)

Error(s) in loading state_dict for SwinUNETR:
size mismatch for out.conv.conv.weight: copying a param with shape torch.Size([14, 12, 1, 1, 1]) from checkpoint, the shape in current model is torch.Size([5, 12, 1, 1, 1]).
size mismatch for out.conv.conv.bias: copying a param with shape torch.Size([14]) from checkpoint, the shape in current model is torch.Size([5]).

[tangy5]

Hi @MasalaKimchi thanks for the question.
You are using the BTCV specific model which is trained with 14 channels output. There might be two options:

1. If you do want to use this model weights, you could load all parameter weights but leave the last segmentation head layers from scratch.
2. If you are following the research-contribution swin-unetr repo. https://github.com/Project-MONAI/research-contributions/tree/main/SwinUNETR/BTCV#models The framework is using the pre-trianed weights with only the encoder which is this model (https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/model_swinvit.pt), not the fine-tunned model.
   Using the SSL pretrained weights, you can train your task with any output channels.

Thanks!