search

sct-pipeline / contrast-agnostic-softseg-spinalcord

Contrast-agnostic spinal cord segmentation project with softseg
MIT License
4 stars 3 forks source link

Adding feature to load weights from pre-trained checkpoint #85

Open rohanbanerjee opened 8 months ago

rohanbanerjee commented 8 months ago

Related issue #83 .

What is the PR?

The current training script does not have the functionality to train model by loading pre-trained checkpoints to initialize fine-tuning. The changes in this PR allows users to use the -rfc or --resume-from-checkpoint flag to load the trained model checkpoint and start the training.

To test the PR:

Load the weights from the location of the trained contrast-agnostic model checkpoint: duke/temp/muena/contrast-agnostic/final_monai_model/nnunet_nf=32_DS=1_opt=adam_lr=0.001_AdapW_CCrop_bs=2_64x192x320_20230918-2253/best_dice_model.ckpt and then train a custom model with a train command, for example:

  1. cd monai
  2. run 
    python main.py --model nnunet --enable_DS -initf 32 -me 1000 -bs 4 -opt adam -lr 0.001 --label-type soft -cve 1 -rfc ~/duke/temp/muena/contrast-agnostic/final_monai_model/nnunet_nf=32_DS=1_opt=adam_lr=0.001_AdapW_CCrop_bs=2_64x192x320_20230918-2253/best_model_dice.ckpt

tagging @naga-karthik