uni-medical / STU-Net

The largest pre-trained medical image segmentation model (1.4B parameters) based on the largest public dataset (>100k annotations), up until April 2023.
Apache License 2.0
279 stars 23 forks source link

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

Open Suchcools opened 3 months ago

Suchcools commented 3 months ago

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,

Ziyan-Huang commented 3 months ago

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

Suchcools commented 3 months ago

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