UNETR in MONAI 0.9.0

[fengling0410]

Hi all,

I have noticed that the UNETR model in MONAI 0.9.0 release has been modified compared to 0.8.1. I found that input_rearrange, out_rearrange were deleted in transformer block and the order of conv3 and instance norm has been shifted in encoder block and decoder block (please correct me if I'm wrong). May I ask about the intuition behind these modifications? I believe that will help me a lot to better use this model.

Besides, I have trained a UNETR model for weeks but the pre-trained weights are now incompatible with the latest UNETR structure. Is there a way that I could load the pre-trained weights to the latest model?

Thank you in advance and I sincerely appreciate any help regarding this :)

[ahatamiz]

@yiheng-wang-nv would appreciate if you can please share some insights.

[yiheng-wang-nv]

Hi @fengling0410 , there are some changes after 0.8.1, the mainly purpose is to make UNETR support torchscript, and the network structure is actually keep the same. You can also refer to: https://github.com/Project-MONAI/MONAI/pull/3923/files and https://github.com/Project-MONAI/MONAI/pull/3782/files
As for the conv3 and norm3 of UnetResBlock, they are removed if not used (see: #4314).
For your trained weights, it may need some changes of the state dict in order to use in the latest version. Could you please provide an example here (like replace the weights by random tensors with same shapes), thus I'm able to check how to modify it?

[fengling0410]

Hi @yiheng-wang-nv, Thank you so much for your reply. I will go through the link you've provided in details. For the pre-trained weights, I'm not sure how I could attach them here since it is very long. I print out the keys and corresponding weights shape of the state_dict and hope it helps :)

vit.patch_embedding.position_embeddings torch.Size([1, 216, 768]) vit.patch_embedding.cls_token torch.Size([1, 1, 768]) vit.patch_embedding.patch_embeddings.1.weight torch.Size([768, 4096]) vit.patch_embedding.patch_embeddings.1.bias torch.Size([768]) vit.blocks.0.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.0.mlp.linear1.bias torch.Size([3072]) vit.blocks.0.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.0.mlp.linear2.bias torch.Size([768]) vit.blocks.0.norm1.weight torch.Size([768]) vit.blocks.0.norm1.bias torch.Size([768]) vit.blocks.0.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.0.attn.out_proj.bias torch.Size([768]) vit.blocks.0.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.0.norm2.weight torch.Size([768]) vit.blocks.0.norm2.bias torch.Size([768]) vit.blocks.1.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.1.mlp.linear1.bias torch.Size([3072]) vit.blocks.1.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.1.mlp.linear2.bias torch.Size([768]) vit.blocks.1.norm1.weight torch.Size([768]) vit.blocks.1.norm1.bias torch.Size([768]) vit.blocks.1.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.1.attn.out_proj.bias torch.Size([768]) vit.blocks.1.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.1.norm2.weight torch.Size([768]) vit.blocks.1.norm2.bias torch.Size([768]) vit.blocks.2.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.2.mlp.linear1.bias torch.Size([3072]) vit.blocks.2.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.2.mlp.linear2.bias torch.Size([768]) vit.blocks.2.norm1.weight torch.Size([768]) vit.blocks.2.norm1.bias torch.Size([768]) vit.blocks.2.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.2.attn.out_proj.bias torch.Size([768]) vit.blocks.2.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.2.norm2.weight torch.Size([768]) vit.blocks.2.norm2.bias torch.Size([768]) vit.blocks.3.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.3.mlp.linear1.bias torch.Size([3072]) vit.blocks.3.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.3.mlp.linear2.bias torch.Size([768]) vit.blocks.3.norm1.weight torch.Size([768]) vit.blocks.3.norm1.bias torch.Size([768]) vit.blocks.3.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.3.attn.out_proj.bias torch.Size([768]) vit.blocks.3.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.3.norm2.weight torch.Size([768]) vit.blocks.3.norm2.bias torch.Size([768]) vit.blocks.4.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.4.mlp.linear1.bias torch.Size([3072]) vit.blocks.4.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.4.mlp.linear2.bias torch.Size([768]) vit.blocks.4.norm1.weight torch.Size([768]) vit.blocks.4.norm1.bias torch.Size([768]) vit.blocks.4.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.4.attn.out_proj.bias torch.Size([768]) vit.blocks.4.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.4.norm2.weight torch.Size([768]) vit.blocks.4.norm2.bias torch.Size([768]) vit.blocks.5.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.5.mlp.linear1.bias torch.Size([3072]) vit.blocks.5.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.5.mlp.linear2.bias torch.Size([768]) vit.blocks.5.norm1.weight torch.Size([768]) vit.blocks.5.norm1.bias torch.Size([768]) vit.blocks.5.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.5.attn.out_proj.bias torch.Size([768]) vit.blocks.5.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.5.norm2.weight torch.Size([768]) vit.blocks.5.norm2.bias torch.Size([768]) vit.blocks.6.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.6.mlp.linear1.bias torch.Size([3072]) vit.blocks.6.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.6.mlp.linear2.bias torch.Size([768]) vit.blocks.6.norm1.weight torch.Size([768]) vit.blocks.6.norm1.bias torch.Size([768]) vit.blocks.6.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.6.attn.out_proj.bias torch.Size([768]) vit.blocks.6.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.6.norm2.weight torch.Size([768]) vit.blocks.6.norm2.bias torch.Size([768]) vit.blocks.7.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.7.mlp.linear1.bias torch.Size([3072]) vit.blocks.7.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.7.mlp.linear2.bias torch.Size([768]) vit.blocks.7.norm1.weight torch.Size([768]) vit.blocks.7.norm1.bias torch.Size([768]) vit.blocks.7.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.7.attn.out_proj.bias torch.Size([768]) vit.blocks.7.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.7.norm2.weight torch.Size([768]) vit.blocks.7.norm2.bias torch.Size([768]) vit.blocks.8.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.8.mlp.linear1.bias torch.Size([3072]) vit.blocks.8.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.8.mlp.linear2.bias torch.Size([768]) vit.blocks.8.norm1.weight torch.Size([768]) vit.blocks.8.norm1.bias torch.Size([768]) vit.blocks.8.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.8.attn.out_proj.bias torch.Size([768]) vit.blocks.8.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.8.norm2.weight torch.Size([768]) vit.blocks.8.norm2.bias torch.Size([768]) vit.blocks.9.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.9.mlp.linear1.bias torch.Size([3072]) vit.blocks.9.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.9.mlp.linear2.bias torch.Size([768]) vit.blocks.9.norm1.weight torch.Size([768]) vit.blocks.9.norm1.bias torch.Size([768]) vit.blocks.9.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.9.attn.out_proj.bias torch.Size([768]) vit.blocks.9.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.9.norm2.weight torch.Size([768]) vit.blocks.9.norm2.bias torch.Size([768]) vit.blocks.10.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.10.mlp.linear1.bias torch.Size([3072]) vit.blocks.10.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.10.mlp.linear2.bias torch.Size([768]) vit.blocks.10.norm1.weight torch.Size([768]) vit.blocks.10.norm1.bias torch.Size([768]) vit.blocks.10.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.10.attn.out_proj.bias torch.Size([768]) vit.blocks.10.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.10.norm2.weight torch.Size([768]) vit.blocks.10.norm2.bias torch.Size([768]) vit.blocks.11.mlp.linear1.weight torch.Size([3072, 768]) vit.blocks.11.mlp.linear1.bias torch.Size([3072]) vit.blocks.11.mlp.linear2.weight torch.Size([768, 3072]) vit.blocks.11.mlp.linear2.bias torch.Size([768]) vit.blocks.11.norm1.weight torch.Size([768]) vit.blocks.11.norm1.bias torch.Size([768]) vit.blocks.11.attn.out_proj.weight torch.Size([768, 768]) vit.blocks.11.attn.out_proj.bias torch.Size([768]) vit.blocks.11.attn.qkv.weight torch.Size([2304, 768]) vit.blocks.11.norm2.weight torch.Size([768]) vit.blocks.11.norm2.bias torch.Size([768]) vit.norm.weight torch.Size([768]) vit.norm.bias torch.Size([768]) encoder1.layer.conv1.conv.weight torch.Size([16, 1, 3, 3, 3]) encoder1.layer.conv2.conv.weight torch.Size([16, 16, 3, 3, 3]) encoder1.layer.conv3.conv.weight torch.Size([16, 1, 1, 1, 1]) encoder2.transp_conv_init.conv.weight torch.Size([768, 32, 2, 2, 2]) encoder2.blocks.0.0.conv.weight torch.Size([32, 32, 2, 2, 2]) encoder2.blocks.0.1.conv1.conv.weight torch.Size([32, 32, 3, 3, 3]) encoder2.blocks.0.1.conv2.conv.weight torch.Size([32, 32, 3, 3, 3]) encoder2.blocks.0.1.conv3.conv.weight torch.Size([32, 32, 1, 1, 1]) encoder2.blocks.1.0.conv.weight torch.Size([32, 32, 2, 2, 2]) encoder2.blocks.1.1.conv1.conv.weight torch.Size([32, 32, 3, 3, 3]) encoder2.blocks.1.1.conv2.conv.weight torch.Size([32, 32, 3, 3, 3]) encoder2.blocks.1.1.conv3.conv.weight torch.Size([32, 32, 1, 1, 1]) encoder3.transp_conv_init.conv.weight torch.Size([768, 64, 2, 2, 2]) encoder3.blocks.0.0.conv.weight torch.Size([64, 64, 2, 2, 2]) encoder3.blocks.0.1.conv1.conv.weight torch.Size([64, 64, 3, 3, 3]) encoder3.blocks.0.1.conv2.conv.weight torch.Size([64, 64, 3, 3, 3]) encoder3.blocks.0.1.conv3.conv.weight torch.Size([64, 64, 1, 1, 1]) encoder4.transp_conv_init.conv.weight torch.Size([768, 128, 2, 2, 2]) decoder5.transp_conv.conv.weight torch.Size([768, 128, 2, 2, 2]) decoder5.conv_block.conv1.conv.weight torch.Size([128, 256, 3, 3, 3]) decoder5.conv_block.conv2.conv.weight torch.Size([128, 128, 3, 3, 3]) decoder5.conv_block.conv3.conv.weight torch.Size([128, 256, 1, 1, 1]) decoder4.transp_conv.conv.weight torch.Size([128, 64, 2, 2, 2]) decoder4.conv_block.conv1.conv.weight torch.Size([64, 128, 3, 3, 3]) decoder4.conv_block.conv2.conv.weight torch.Size([64, 64, 3, 3, 3]) decoder4.conv_block.conv3.conv.weight torch.Size([64, 128, 1, 1, 1]) decoder3.transp_conv.conv.weight torch.Size([64, 32, 2, 2, 2]) decoder3.conv_block.conv1.conv.weight torch.Size([32, 64, 3, 3, 3]) decoder3.conv_block.conv2.conv.weight torch.Size([32, 32, 3, 3, 3]) decoder3.conv_block.conv3.conv.weight torch.Size([32, 64, 1, 1, 1]) decoder2.transp_conv.conv.weight torch.Size([32, 16, 2, 2, 2]) decoder2.conv_block.conv1.conv.weight torch.Size([16, 32, 3, 3, 3]) decoder2.conv_block.conv2.conv.weight torch.Size([16, 16, 3, 3, 3]) decoder2.conv_block.conv3.conv.weight torch.Size([16, 32, 1, 1, 1]) out.conv.conv.weight torch.Size([2, 16, 1, 1, 1]) out.conv.conv.bias torch.Size([2])

And here is the error output when I was trying to load the weights into 0.9.0 UNETR model:

[yiheng-wang-nv]

Hi @fengling0410 , could you please try to remove these keys, and compare the consistency? Or if you can provide the hyperparameters of the network, I'm able to help to reproduce the issue.

[fengling0410]

Hi @yiheng-wang-nv, thank you for your insights and sorry about my late reply. It works if I remove the corresponding weights in the pre-trained state_dict, and update the state_dict of the new model. As far as I understand, since the conv3 in UNETR encoder seems to be useless, thus deleting weights of conv3 doesn't harm the model performance. May I ask if my understanding is correct or not?

And I'm still a little confused about the "vit.patch_embedding.cls_token". According to my understanding, the cls_token is included in hidden_state only if the "classification" argument is set to be true in ViT, but the ViT backbone in UNETR has classification=false. Therefore I'm not quite sure why there is "vit.patch_embedding.cls_token" weight in UNETR of 0.8.1 version. Could you please give some insights about this?

Here is the code snippet I mentioned above in ViT source code:
if self.classification: self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size)) self.classification_head = nn.Sequential(nn.Linear(hidden_size, num_classes), nn.Tanh())

Thank you again and I really appreciate any insights about this :)

[fengling0410]

Hi @yiheng-wang-nv, just to update that deleting the four parameters works and the validation scores is the same as using the UNETR of version 0.8.1.

This is what I did:

pretrainedunetr = torch.load(model_path)
pretrained_dict = {k: v for k, v in pretrainedunetr.items() if k not in ['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']}
model_dict = model.state_dict()
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

[yiheng-wang-nv]

Hi @fengling0410 , cls_token was set during the initial implementation of ViT (see: https://github.com/Project-MONAI/MONAI/pull/2476/files#diff-bf3b32131fd5ac1106b8b5844cebb54d1f86b22665f6ed8fb9dfc7038dd7077eR99), no matter it is used of not.

[ahatamiz]

Hi @yiheng-wang-nv

Thanks for the inputs. Is it possible for your to create a guide on how to use UNETR, if trained on older versions, on 0.9.0 ? I think others may run into other issues.

cc @Nic-Ma

[yiheng-wang-nv]

HI @ahatamiz , thanks for the advice. Except for the guide, actually I'm not sure if it is a good way to use DynUNet blocks in UNETR blocks directly. In the future, if we modify one of the networks, the other one may still get an issue.
Hi @wyli and @Nic-Ma , what do you think?

[ahatamiz]

Hi @yiheng-wang-nv , thanks for the comment. In the paper, we mention that the decoder follows a UNet-type decoder, so would be ideal to use components from DynUNet as much as possible. Currently, the breaking changes are mainly for UNETR patch embedding due to some changes to support Torchscript. It would be great to just highlight these ( which should be minimal).

For DynUNet, do you anticipate breaking changes in the future ? if so, a migration guide that covers both may be sufficient.