diff --git a/Pretraining and Finetuning/End_to_End/main_finetune.py b/Pretraining and Finetuning/End_to_End/main_finetune.py index 31cef11..f77c78f 100644 --- a/Pretraining and Finetuning/End_to_End/main_finetune.py +++ b/Pretraining and Finetuning/End_to_End/main_finetune.py @@ -51,6 +51,7 @@ # ft: continue training parser.add_argument('--ft', type=str, default='False', choices=['True', 'False'], help='finetune model') +parser.add_argument('--load', type=str, default=None, choices=['backbone','network'], help='loaded model part') # must have resume parser.add_argument('--resume', type=str, default=None, help='resume name') @@ -650,4 +651,4 @@ def validation(args, model, valid_loader): torch.save({'epoch': epoch, 'state_dict': model.module.encoder.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler':scheduler.state_dict(), 'loss_finetune':np.array(losses)}, filename) -logger.info('################# Fine tune model save finished! ######################') \ No newline at end of file +logger.info('################# Fine tune model save finished! ######################')