Auto3dSeg Failing for SwinUnet - Incorrect Parameter Dimensions

[peterhessey]

I'm trying to replicate the Auto3dSeg results on the medical decathlon prostate data (task_05). However, I'm running into the following error during the SwinUnet training stage:

RuntimeError: Error(s) in loading state_dict for SwinUNETR:
	size mismatch for swinViT.patch_embed.proj.weight: copying a param with shape torch.Size([48, 1, 2, 2, 2]) from checkpoint, the shape in current model is torch.Size([48, 2, 2, 2, 2]).
	size mismatch for encoder1.layer.conv1.conv.weight: copying a param with shape torch.Size([48, 1, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([48, 2, 3, 3, 3]).
	size mismatch for encoder1.layer.conv3.conv.weight: copying a param with shape torch.Size([48, 1, 1, 1, 1]) from checkpoint, the shape in current model is torch.Size([48, 2, 1, 1, 1]).

Looking into the hyperparameters.yaml shows that the input_channels is being set to 2, which explains the above error as it appears to be attempting to load a 1-channel input layer. Changing this value to 1 allows the model to load the state dict correctly, but then fails during training as the prostate data has two channels.

I didn't have this error ~2 months ago when running on the prostate data.

Any suggestions on what I can do to fix this?

Related issue opened here

[dongyang0122]

hi @tangy5, can you look into this ticket? I am thinking it is possible that loading pre-trained weights does not support multi-channel input. Thanks in advance!

[dongyang0122]

add @heyufan1995 here for further discussion.

[heyufan1995]

@peterhessey Hi, it's a bug in loading pretrained weights. Once the scripts are generated, please change the line 192 in scripts.train from "store_dict[key] = model_dict[key]" to "store_dict[key].copy_(model_dict[key])". We will fix this in the algorithm template.

[tangy5]

@heyufan1995 @dongyang0122 , sorry I didn't see this in an earlier time. I saw Yufan figured out the problem, thank you.
@yufan Will you submit a PR to fix this or I can help create a PR?

[peterhessey]

Awesome, thank you @heyufan1995 !