Finetuning of multiTalent multistem model

[plbenveniste]

This issue is related to this #53.
The objective is to fine-tune the multiTalent multistem model, which was pretrained on MRI images.
Pretrained model is stored here: https://zenodo.org/records/13753413
Code of the MultiTalent repo: https://github.com/MIC-DKFZ/MultiTalent

[plbenveniste]

I was able (with help) to fix the code for fine-tuning the model: plbenveniste/nnUNet#3

There are a lof of possibilities for the choice of the stem: MIC-DKFZ/MultiTalent#19 (comment)

[plbenveniste]

Some interesting possibilities are:

• 052,CT,"background, Lungtumor",Dataset052_StructSeg2019_Task4_Lung_GTV,https://structseg2019.grand-challenge.org/Dataset/
• 051,CT,"background, left lung, right lung, heart, esophagus, bronchies, spinal cord nerve thingy",Dataset051_StructSeg2019_Task3_Thoracic_OAR,https://structseg2019.grand-challenge.org/Dataset/
• 351_5,CTzscore,"background, vertebrae_C4, vertebrae_C5, vertebrae_C6, vertebrae_C7, vertebrae_L1, vertebrae_L2, vertebrae_L3, vertebrae_L4, vertebrae_L5, vertebrae_S1, vertebrae_T1, vertebrae_T2, vertebrae_T3, vertebrae_T4, vertebrae_T5, vertebrae_T6, vertebrae_T7, vertebrae_T8, vertebrae_T9",Dataset351_TotalSegmentator_v2,
• 011,"DWI, FLAIR, T1, T2","background, stroke",Dataset011_ISLES,ISCHEMIC STROKE LESION SEGMENTATION
• 201,Flair,"background, MSlesion",Dataset201_MS_Flair,MS-lesion
• 351_4,CTzscore,"background, rib_right_12, sacrum, scapula_left, scapula_right, skull, small_bowel, spinal_cord, spleen, sternum, stomach, subclavian_artery_left, subclavian_artery_right, superior_vena_cava, thyroid_gland, trachea, urinary_bladder, vertebrae_C1, vertebrae_C2, vertebrae_C3",Dataset351_TotalSegmentator_v2,
• 035,"flair, mprage, pd, t2","background, lesion",Dataset035_ISBILesionSegmentation,nothing

I'll focus in priority on 035 (I like that it is MS lesion segmentation with 4 contrasts) and 351_5 (as it covers a large portion of the spinal cord).

[plbenveniste]

I have the following issue when finetuning the model on compute Canada for jobs 250 and 260. These jobs correspond to stem_ID 035.

Traceback (most recent call last):
  File "/home/p/plb/links/projects/aip-jcohen/plb/final_trainings/nnUNet/nnunetv2/run/run_training_MT.py", line 283, in <module>
    run_training_entry()
  File "/home/p/plb/links/projects/aip-jcohen/plb/final_trainings/nnUNet/nnunetv2/run/run_training_MT.py", line 271, in run_training_entry
    run_training(args.dataset_name_or_id, args.configuration, args.fold, args.tr, args.p, args.pretrained_weights,
  File "/home/p/plb/links/projects/aip-jcohen/plb/final_trainings/nnUNet/nnunetv2/run/run_training_MT.py", line 201, in run_training
    maybe_load_checkpoint(nnunet_trainer, continue_training, only_run_validation, pretrained_weights, stemid=stemid)
  File "/home/p/plb/links/projects/aip-jcohen/plb/final_trainings/nnUNet/nnunetv2/run/run_training_MT.py", line 94, in maybe_load_checkpoint
    load_pretrained_weights_MT(nnunet_trainer.network, pretrained_weights_file, verbose=True, stemid=stemid)
  File "/project/6102268/plb/final_trainings/nnUNet/nnunetv2/run/load_pretrained_weights_MT.py", line 115, in load_pretrained_weights_MT
    mod.load_state_dict(model_dict)
  File "/project/6102268/plb/final_trainings/.venv_job200/lib/python3.11/site-packages/torch/nn/modules/module.py", line 2593, in load_state_dict
    raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for ResidualEncoderUNet:
	size mismatch for encoder.stem.convs.0.conv.weight: copying a param with shape torch.Size([32, 4, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([32, 1, 3, 3, 3]).
	size mismatch for encoder.stem.convs.0.all_modules.0.weight: copying a param with shape torch.Size([32, 4, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([32, 1, 3, 3, 3]).
	size mismatch for decoder.encoder.stem.convs.0.conv.weight: copying a param with shape torch.Size([32, 4, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([32, 1, 3, 3, 3]).
	size mismatch for decoder.encoder.stem.convs.0.all_modules.0.weight: copying a param with shape torch.Size([32, 4, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([32, 1, 3, 3, 3]).

[plbenveniste]

I think that the issue is do to the fact that the input for stem 035 is that it expects 4 images as input when I am only feeding one.

[plbenveniste]

I removed trainings which included stem 035 and only kept stem 351_5.