PuzzleTuning/utils/transfermodel.py

"""
Transfer PuzzleTuning Pre-Training checkpoints    Script  ver： Oct 23rd 17:00

write a model based on the weight of a checkpoint file
EG: create a vit-base based on PuzzleTuning SAE

"""
import argparse

import sys
sys.path.append('..')
import os
import torch
import torch.nn as nn

from Backbone import getmodel, GetPromptModel
from SSL_structures import SAE


# Transfer pretrained MSHT checkpoints to normal model state_dict
def transfer_model_encoder(check_point_path, save_model_path, model_idx='ViT', prompt_mode=None,
                           Prompt_Token_num=20, edge_size=384, given_name=None):
    if not os.path.exists(save_model_path):
        os.makedirs(save_model_path)

    if given_name is not None:
        given_path = os.path.join(save_model_path, given_name)
    else:
        given_path = None

    if prompt_mode == "Deep" or prompt_mode == "Shallow":
        model = GetPromptModel.build_promptmodel(edge_size=edge_size, model_idx=model_idx, patch_size=16,
                                                 Prompt_Token_num=Prompt_Token_num, VPT_type=prompt_mode,
                                                 base_state_dict=None)
    # elif prompt_mode == "Other" or prompt_mode == None:
    else:
        model = getmodel.get_model(model_idx=model_idx, pretrained_backbone=False, edge_size=edge_size)
    '''
    state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
    TempBest_state = {'model': best_model_wts, 'epoch': best_epoch_idx}
    '''
    state = torch.load(check_point_path)

    transfer_name = os.path.splitext(os.path.split(check_point_path)[1])[0] + '_of_'

    try:
        model_state = state['model']
        try:
            print("checkpoint epoch", state['epoch'])
            if prompt_mode is not None:
                save_model_path = os.path.join(save_model_path, transfer_name +
                                               model_idx + '_E_' + str(state['epoch']) + '_promptstate' + '.pth')
            else:
                save_model_path = os.path.join(save_model_path, transfer_name +
                                               model_idx + '_E_' + str(state['epoch']) + '_transfer' + '.pth')

        except:
            print("no 'epoch' in state")
            if prompt_mode is not None:
                save_model_path = os.path.join(save_model_path, transfer_name + model_idx + '_promptstate' + '.pth')
            else:
                save_model_path = os.path.join(save_model_path, transfer_name + model_idx + '_transfer' + '.pth')
    except:
        print("not a checkpoint state (no 'model' in state)")
        model_state = state
        if prompt_mode is not None:
            save_model_path = os.path.join(save_model_path, transfer_name + model_idx + '_promptstate' + '.pth')
        else:
            save_model_path = os.path.join(save_model_path, transfer_name + model_idx + '_transfer' + '.pth')

    try:
        model.load_state_dict(model_state)
        print("model loaded")
        print("model :", model_idx)
        gpu_use = 0
    except:
        try:
            model = nn.DataParallel(model)
            model.load_state_dict(model_state, False)
            print("DataParallel model loaded")
            print("model :", model_idx)
            gpu_use = -1
        except:
            print("model loading erro!!")
            gpu_use = -2

    if given_path is not None:
        save_model_path = given_path

    if gpu_use == -1:
        # print(model)
        if prompt_mode is not None:
            prompt_state_dict = model.module.obtain_prompt()
            # fixme maybe bug at DP module.obtain_prompt, just model.obtain_prompt is enough
            print('prompt obtained')
            torch.save(prompt_state_dict, save_model_path)
        else:
            torch.save(model.module.state_dict(), save_model_path)
        print('model trained by multi-GPUs has its single GPU copy saved at ', save_model_path)

    elif gpu_use == 0:
        if prompt_mode is not None:
            prompt_state_dict = model.obtain_prompt()
            print('prompt obtained')
            torch.save(prompt_state_dict, save_model_path)
        else:
            torch.save(model.state_dict(), save_model_path)
        print('model trained by a single GPU has been saved at ', save_model_path)
    else:
        print('erro')


def transfer_model_decoder(check_point_path, save_model_path,
                           model_idx='sae_vit_base_patch16_decoder', dec_idx='swin_unet',
                           prompt_mode=None, Prompt_Token_num=20, edge_size=384):

    if not os.path.exists(save_model_path):
        os.makedirs(save_model_path)

    state = torch.load(check_point_path)

    transfer_name = os.path.splitext(os.path.split(check_point_path)[1])[0] + '_of_'

    model = SAE.__dict__[model_idx](img_size=edge_size, prompt_mode=prompt_mode, Prompt_Token_num=Prompt_Token_num,
                                    basic_state_dict=None, dec_idx=dec_idx)

    try:
        model_state = state['model']
        try:
            print("checkpoint epoch", state['epoch'])
            save_model_path = os.path.join(save_model_path, transfer_name + 'Decoder_' + dec_idx + '_E_'
                                           + str(state['epoch']) + '.pth')


        except:
            print("no 'epoch' in state")
            save_model_path = os.path.join(save_model_path, transfer_name + 'Decoder_' + dec_idx + '.pth')
    except:
        print("not a checkpoint state (no 'model' in state)")
        model_state = state
        save_model_path = os.path.join(save_model_path, transfer_name + 'Decoder_' + dec_idx + '.pth')

    try:
        model.load_state_dict(model_state)
        print("model loaded")
        print("model :", model_idx)
        gpu_use = 0
    except:
        try:
            model = nn.DataParallel(model)
            model.load_state_dict(model_state, False)
            print("DataParallel model loaded")
            print("model :", model_idx)
            gpu_use = -1
        except:
            print("model loading erro!!")
            gpu_use = -2

    else:
        model = model.decoder

    if gpu_use == -1:
        torch.save(model.module.decoder.state_dict(), save_model_path)
        print('model trained by multi-GPUs has its single GPU copy saved at ', save_model_path)

    elif gpu_use == 0:
        torch.save(model.state_dict(), save_model_path)
        print('model trained by a single GPU has been saved at ', save_model_path)
    else:
        print('erro')


def get_args_parser():
    parser = argparse.ArgumentParser('Take pre-trained model from PuzzleTuning', add_help=False)

    # Model Name or index
    parser.add_argument('--given_name', default=None, type=str, help='name of the weight-state-dict')
    parser.add_argument('--model_idx', default='ViT', type=str, help='taking the weight to the specified model')
    parser.add_argument('--edge_size', default=224, type=int, help='images input size for model')

    # PromptTuning
    parser.add_argument('--PromptTuning', default=None, type=str,
                        help='Deep/Shallow to use Prompt Tuning model instead of Finetuning model, by default None')
    # Prompt_Token_num
    parser.add_argument('--Prompt_Token_num', default=20, type=int, help='Prompt_Token_num')

    # PATH settings
    parser.add_argument('--checkpoint_path', default=None, type=str, help='check_point_path')
    parser.add_argument('--save_model_path', default=None, type=str, help='out put weight path for pre-trained model')

    return parser


def main(args):
    # fixme: now need a CUDA device as the model is save as a CUDA model!


    # PuzzleTuning Template
    """
    # Prompt
    # transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep', Prompt_Token_num=20,given_name=given_name)

    # not prompt model
    # transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, given_name=given_name)

    # decoder
    # transfer_model_decoder(checkpoint_path, save_model_path, model_idx='sae_vit_base_patch16_decoder', dec_idx='swin_unet', edge_size=224, prompt_mode='Deep')


    # PuzzleTuning Experiments transfer records:
    # 1 周期puzzle自动减小ratio，自动loop变化size 迁移timm，PromptTuning：VPT-Deep，seg_decoder：None (核心方法)
    # ViT_b16_224_timm_PuzzleTuning_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_vit_base_patch16_Prompt_Deep_tokennum_20_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_SAE_CPIAm_Prompt_Deep_tokennum_20_E_50_promptstate.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep',
                           Prompt_Token_num=20,given_name=given_name)

    # PuzzleTuning Ablation studies:SAE+不同curriculum+不同VPT/ViT
    # 2 周期puzzle自动减小ratio，自动loop变化size 迁移timm，PromptTuning：None，seg_decoder：None
    # ViT_b16_224_timm_PuzzleTuning_SAE_CPIAm_E_199.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_vit_base_patch16_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_SAE_CPIAm_E_199.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, given_name=given_name)

    # 3 固定puzzle ratio，固定patch size 迁移timm，PromptTuning：VPT-Deep，seg_decoder：None  (服务器pt1)
    # ViT_b16_224_timm_PuzzleTuning_fixp16fixr25_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_fixp16fixr25_vit_base_Prompt_Deep_tokennum_20_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_SAE_fixp16fixr25_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep',
                           Prompt_Token_num=20, given_name=given_name)

    # 4 固定puzzle ratio，固定patch size 迁移timm，PromptTuning：None，seg_decoder：None (服务器pt2)
    # ViT_b16_224_timm_PuzzleTuning_fixp16fixr25_SAE_CPIAm_E_199.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_fixp16fixr25_vit_base_patch16_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_SAE_fixp16fixr25_CPIAm_E_199.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, given_name=given_name)

    # 5 变化puzzle ratio，固定patch size 迁移timm，PromptTuning：VPT-Deep，seg_decoder：None, strategy: ratio-decay (服务器pt3)
    # ViT_b16_224_timm_PuzzleTuning_fixp16ratiodecay_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_fixp16ratiodecay_vit_base_Prompt_Deep_tokennum_20_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_SAE_fixp16ratiodecay_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep',
                           Prompt_Token_num=20, given_name=given_name)

    # 6 变化puzzle ratio，固定patch size 迁移timm，PromptTuning：None，seg_decoder：None (服务器pt4)
    # ViT_b16_224_timm_PuzzleTuning_fixp16ratiodecay_SAE_CPIAm_E_199.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_fixp16ratiodecay_vit_base_patch16_tr_timm_CPIAm/PuzzleTuning_sae_vit_base_patch16_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_SAE_fixp16ratiodecay_CPIAm_E_199.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, given_name=given_name)

    # PuzzleTuning Ablation studies:上游不要puzzle 所以是 VPT+MAE
    # 7 MAE+VPT，迁移timm，PromptTuning：VPT-Deep，seg_decoder：None (A40*4服务器pt5)
    # ViT_b16_224_timm_PuzzleTuning_MAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_MAE_vit_base_Prompt_Deep_tokennum_20_tr_timm_CPIAm/PuzzleTuning_mae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_timm_PuzzleTuning_MAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep',
                           Prompt_Token_num=20, given_name=given_name)

    # 8 周期puzzle自动减小ratio，自动loop变化size 迁移MAEImageNet，PromptTuning：VPT-Deep，seg_decoder：None (A100-PCIE*2 服务器pt6)
    # ViT_b16_224_MAEImageNet_PuzzleTuning_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_vit_base_patch16_Prompt_Deep_tokennum_20_tr_MAEImageNet_CPIAm/PuzzleTuning_sae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_MAEImageNet_PuzzleTuning_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep',
                           Prompt_Token_num=20, given_name=given_name)

    # 9 周期puzzle自动减小ratio，自动loop变化size 迁移Random，PromptTuning：VPT-Deep，seg_decoder：None  (A40*4服务器pt7)
    # ViT_b16_224_Random_PuzzleTuning_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_vit_base_patch16_Prompt_Deep_tokennum_20_tr_Random_CPIAm/PuzzleTuning_sae_vit_base_patch16_Prompt_Deep_tokennum_20_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_Random_PuzzleTuning_SAE_CPIAm_Prompt_Deep_tokennum_20_E_199_promptstate.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, prompt_mode='Deep',
                           Prompt_Token_num=20, given_name=given_name)
                           
    # 10 周期puzzle自动减小ratio，自动loop变化size 迁移Random，PromptTuning：None，seg_decoder：None  (4090*6服务器pt8)
    # ViT_b16_224_MAEImageNet_PuzzleTuning_SAE_CPIAm_E_199.pth
    checkpoint_path = '/root/autodl-tmp/runs/PuzzleTuning_SAE_vit_base_patch16_tr_MAEImageNet_CPIAm/PuzzleTuning_sae_vit_base_patch16_checkpoint-199.pth'
    save_model_path = '/root/autodl-tmp/output_models'
    given_name = r'ViT_b16_224_MAEImageNet_PuzzleTuning_SAE_CPIAm_E_199.pth'
    transfer_model_encoder(checkpoint_path, save_model_path, model_idx='ViT', edge_size=224, given_name=given_name)
    """

    transfer_model_encoder(args.checkpoint_path, args.save_model_path,
                           model_idx=args.model_idx, edge_size=args.edge_size,
                           prompt_mode=args.PromptTuning, Prompt_Token_num=args.Prompt_Token_num,
                           given_name=args.given_name)


if __name__ == '__main__':
    args = get_args_parser()
    args = args.parse_args()

    main(args)