Fine-tuning and transfer learning

[ForgottenOneNyx]

Hello.

Thank you so much for your work on this project, it is very amazing.

I do have a question regarding how a nnUNet model (that I have trained for my own task) can be fine tuned on a new dataset. I am particularly interested in two scenarios.

1. When I want to fine tune a trained model (for example: which takes a 2 modality input and is a 4 class multi-class segmentation) to a new model which again takes a 2 modality input and is a 4 class multi-class segmentation. Its just that while fine-tuning it trains using a new dataset (same characteristics as the data the trained model used.)
2. When I want to fine tune a trained model (for example: which takes a 2 modality input and is a 4 class multi-class segmentation) to a new model which take a 2 modality input but is now a 3 class multi-class segmentation. So while fine-tuning it trains using a new dataset (different characteristics as the data the trained model used, because its a 3 class segmentation)

I won't be performing any cross organ segmentations. So its in the same organ space and image space.

Let me know how I can perform the 2 scenarios. I will be happy to answer any queries you have. Thank you in advance.

[FabianIsensee]

Hi, you can do both with nnU-Net. Have a look at nnUNet_train -h for how to load pretrained parameters. It will discard the segmentation layers automatically. You should also implement a custom trainer class with a proper fine tuning scheme so that you don't destroy the pretrained weights at the start of the training. Important: the model you load for pretrained weights MUST have been created using the same plans. You can do this by using -overwrite_plans in nnUNet_plan_and_preprocess Best, Fabian

[Jamshidhsp]

Hi, I am of the opinion to do the same as you mentioned. Could you manage to work it out? @ForgottenOneNyx

[adwaykanhere]

Hi @FabianIsensee

Thanks so much for your work on this amazing tool!

I'm trying to fine-tune the pretrained model for Task006_Lung on a custom TCIA dataset and didn't understand what you meant by

You should also implement a custom trainer class with a proper fine tuning scheme so that you don't destroy the pretrained weights at the start of the training.

I'm using a custom dataset (Task 963) and the following code to run the preprocessing and fine-tuning

!nnUNet_plan_and_preprocess -t 963

!nnUNet_install_pretrained_model_from_zip $PATH_TO_MODEL_FILE

!nnUNet_train 3d_fullres nnUNetTrainerV2 963 0 -pretrained_weights PRETRAINED_WEIGHTS

Is this the right way to run the code? Also I'm not sure if after installing the pre-trained model, will nnUnet automatically take the PRETRAINED_WEIGHTS parameters?

Thanks again for your support!

Adway

[ForgottenOneNyx]

@Jamshidhsp, yes it works for me. Following is the way you can achieve it.

Firstly have the latest version of nnUNet. (At the time it was 1.7.0, I used that)

nnunet - 1.7.0

Idea is to preprocess the "NEW" data on the nnUNet plans for the model you want start fine tuning with.

For example:

Task111_Example2 should be preprocessed using the plans for Task106_Example1

nnUNet_plan_and_preprocess -t YYY -pl3d ExperimentPlanner3D_v21_Pretrained -overwrite_plans /location/to/nnUNet_preprocessed/TaskXXX/nnUNetPlansv2.1_plans_3D.pkl -overwrite_plans_identifier IDENTIFIER

So for this case an example would be:

nnUNet_plan_and_preprocess -t 111 -pl3d ExperimentPlanner3D_v21_Pretrained -overwrite_plans /media/rnd/nnUNet_preprocessed/Task106_Example1/nnUNetPlansv2.1_plans_3D.pkl -overwrite_plans_identifier TRANSFER_TRIAL

And now that we have processed the data into the version of the model we need to fine tune. We just need to train. Which is a little different but similar:

CUDA_VISIBLE_DEVICES=0 nnUNet_train 3d_fullres nnUNetTrainerV2 TaskYYY_ExpName all -p nnUNetPlans_pretrained_IDENTIFIER --npz -pretrained_weights /location/to/model/file/3d_fullres/TaskXXX_FineTuneModelOrigin/nnUNetTrainerV2__nnUNetPlansv2.1/all/model_final_checkpoint.model

So for this case an example would be:

CUDA_VISIBLE_DEVICES=0 nnUNet_train 3d_fullres nnUNetTrainerV2 Task111_Example2 all -p nnUNetPlans_pretrained_TRANSFER_TRIAL --npz -pretrained_weights /media/rnd/nnUNet/RESULTS_FOLDER/nnUNet/3d_fullres/Task106_Example1/nnUNetTrainerV2__nnUNetPlansv2.1/all/model_final_checkpoint.model

[Jamshidhsp]

@ForgottenOneNyx Thank you very much for your explicit explanations. It helps a lot. Actually, I did the same and I got the results. But I think it doesnt's keep the weights and retrains the whole network. Could you find a way to keep some layers retained and only finetune the segment layers?

[adwaykanhere]

@ForgottenOneNyx

Thanks for your post. I'm installing a pretrained model from a zip file and when I run

!nnUNet_install_pretrained_model_from_zip $PATH_TO_MODEL_FILE

I don't get a planner file after I install the pretrained model.

How did you get the file? - /media/rnd/nnUNet_preprocessed/Task106_Example1/nnUNetPlansv2.1_plans_3D.pkl

[ForgottenOneNyx]

@Jamshidhsp Let me look into it today. I shall tell you by tomorrow.

@adwaykanhere I will try to install it by zip. I have never done it. But I shall try and let you know. Give me a day's time

[adwaykanhere]

Thanks @ForgottenOneNyx

If it helps - This is the code that will download and install the pretrained model that i used-

# this will usually take between one and five minutes
!mkdir -p /content/tutorial/models

seg_model_url = "https://www.dropbox.com/s/4ajcm2rvd8z40f1/Task006_Lung.zip?dl=0"
output_path = "/content/tutorial/models/Task006_Lung.zip"
!wget -O $output_path $seg_model_url

!nnUNet_install_pretrained_model_from_zip $output_path

[adwaykanhere]

@ForgottenOneNyx Were you able to figure out if we can load from a zip file?

[mwhchan8]

@ForgottenOneNyx and anyone else: Is there a name for the technique used in the 1st scenario of @ForgottenOneNyx's original post? It is considered transfer learning? What is the exact name of the transfer learning technique then? Is it just fine-tuning? (I'm new to medical segmentation and deep learning. Sorry if the question is dumb)

The 1st scenario:

1. When I want to fine tune a trained model (for example: which takes a 2 modality input and is a 4 class multi-class segmentation) to a new model which again takes a 2 modality input and is a 4 class multi-class segmentation. Its just that while fine-tuning it trains using a new dataset (same characteristics as the data the trained model used.)

[FabianIsensee]

Hard to say. From my perspective the borders between transfer learning and fine tuning are blurry. I would call the procedure here transfer learning and the technique used to do the prodecure fine tuning

[mwhchan8]

Thank you very much for the explanation! @FabianIsensee

[nohalloween]

@ForgottenOneNyx Hi, have you got a solution of the scenarios 2?

[mehnaz1985]

@Jamshidhsp, yes it works for me. Following is the way you can achieve it.

Firstly have the latest version of nnUNet. (At the time it was 1.7.0, I used that)

nnunet - 1.7.0

Idea is to preprocess the "NEW" data on the nnUNet plans for the model you want start fine tuning with.

For example:

Task111_Example2 should be preprocessed using the plans for Task106_Example1

nnUNet_plan_and_preprocess -t YYY -pl3d ExperimentPlanner3D_v21_Pretrained -overwrite_plans /location/to/nnUNet_preprocessed/TaskXXX/nnUNetPlansv2.1_plans_3D.pkl -overwrite_plans_identifier IDENTIFIER

So for this case an example would be:

nnUNet_plan_and_preprocess -t 111 -pl3d ExperimentPlanner3D_v21_Pretrained -overwrite_plans /media/rnd/nnUNet_preprocessed/Task106_Example1/nnUNetPlansv2.1_plans_3D.pkl -overwrite_plans_identifier TRANSFER_TRIAL

And now that we have processed the data into the version of the model we need to fine tune. We just need to train. Which is a little different but similar:

CUDA_VISIBLE_DEVICES=0 nnUNet_train 3d_fullres nnUNetTrainerV2 TaskYYY_ExpName all -p nnUNetPlans_pretrained_IDENTIFIER --npz -pretrained_weights /location/to/model/file/3d_fullres/TaskXXX_FineTuneModelOrigin/nnUNetTrainerV2__nnUNetPlansv2.1/all/model_final_checkpoint.model

So for this case an example would be:

CUDA_VISIBLE_DEVICES=0 nnUNet_train 3d_fullres nnUNetTrainerV2 Task111_Example2 all -p nnUNetPlans_pretrained_TRANSFER_TRIAL --npz -pretrained_weights /media/rnd/nnUNet/RESULTS_FOLDER/nnUNet/3d_fullres/Task106_Example1/nnUNetTrainerV2__nnUNetPlansv2.1/all/model_final_checkpoint.model

could u pls provide the command for the prediction after that?