uni-medical/STU-Net

Fine-tuning models with nnUNetv2 found that model weights are incompatible

Opened this issue · 2 comments

I am encountering an issue while attempting to fine-tune nnUNetv2 with the BRATS21 dataset. Here are the details of my setup and the problem I'm facing:

Dataset: BRATS21
Modalities: T1, T2, T1C, and FLAIR
Data Preparation: I have processed the data following the nnUnetv2 methodology.
Pretrained Weights: I downloaded the base_ep4k.model weights.
Command Used:

python run_finetuning_stunet.py Dataset137 3d_fullres 1 -pretrained_weights /home/linjiawei1/ProjectHub/2024/STU-Net-main/plan_files/base_ep4k.model

Upon running the fine-tuning script, I received a KeyError on line 39 of run_finetuning_stunet.py:

num_inputs = model_dict['conv_blocks_context.0.0.conv1.weight'].shape[1]

The error message is:

KeyError: 'conv_blocks_context.0.0.conv1.weight'

I believe this issue may be due to a mismatch between the expected model architecture and the actual architecture of the base_ep4k.model weights I am using for fine-tuning.

I would appreciate any guidance on how to resolve this issue. Thank you for your time and assistance.
Best regards,

The correct command should include specifying the trainer for fine-tuning. Otherwise, the nnUNet model will be used instead of the STUNet model, and their structures have some differences. Here's the updated command:

python run_finetuning_stunet.py Dataset137 3d_fullres 1 -pretrained_weights /home/linjiawei1/ProjectHub/2024/STU-Net-main/plan_files/base_ep4k.model -tr STUNetTrainer_base_ft

Please try running this updated command and see if it resolves the issue.

Best regards,

Ziyan Huang

It's useful. I don't know much about the configuration and usage of nnUNETv2. Thank you for your contribution and guidance.
Best regards,