model/src/mrl_te_optimization/script/main_train.py

import os,sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
import argparse

parser = argparse.ArgumentParser('the main to train model')
parser.add_argument('--config_file',type=str,required=True)
parser.add_argument('--cuda',type=int,default=None,required=False)
parser.add_argument("--kfold_index",type=int,default=None,required=False)
args = parser.parse_args()

cuda_id = args.cuda
os.environ["CUDA_VISIBLE_DEVICES"] = str(cuda_id)


import time
import torch
import utils
from torch import optim
import numpy as np 
from models import Modules,reader,train_val
from models.ScheduleOptimizer import ScheduledOptim,scheduleoptim_dict_str
from models.popen import Auto_popen
from models.loss import Dynamic_Task_Priority,Dynamic_Weight_Averaging


POPEN = Auto_popen(args.config_file)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.set_num_interop_threads(4)
POPEN.cuda_id = device


POPEN.kfold_index = args.kfold_index
if POPEN.kfold_cv:
    if args.kfold_index is None:
        raise NotImplementedError("please specify the kfold index to perform K fold cross validation")
    POPEN.vae_log_path = POPEN.vae_log_path.replace(".log","_cv%d.log"%args.kfold_index)
    #POPEN.vae_pth_path = POPEN.vae_pth_path.replace(".pth","_cv%d.pth"%args.kfold_index)
    

# Run name
if POPEN.run_name is None:
    run_name = POPEN.model_type + time.strftime("__%Y_%m_%d_%H:%M")
else:
    run_name = POPEN.run_name
    
# log dir
logger = utils.setup_logs(POPEN.vae_log_path)
logger.info(f"  	 	 ==============<<< device used: {device}:{cuda_id}  >>>============== 	 	 ")
#  built model dir or check resume 
POPEN.check_experiment(logger)
#                               |=====================================|
#                               |===========   setup  part  ==========|
#                               |=====================================|
# read data
train_loader,val_loader,test_loader = reader.get_dataloader(POPEN)
# ===========  setup model  ===========
# train_iter = iter(train_loader)
# X,Y  = next(train_iter)
# -- pretrain -- 
if POPEN.pretrain_pth is not None:
    # load pretran model
    pretrain_popen = Auto_popen(os.path.join(utils.script_dir, POPEN.pretrain_pth))
    pretrain_model = pretrain_popen.Model_Class(*pretrain_popen.model_args)

    if not os.path.exists(pretrain_popen.vae_pth_path):
        if type(args.kfold_index) == int:
            pretrain_popen.kfold_index = args.kfold_index
    pretrain_model = utils.load_model(pretrain_popen,pretrain_model,logger)
    

    
    if POPEN.Model_Class == pretrain_popen.Model_Class:
        # if not POPEN.Resumable:
        #     # we only load pre-train for the first time 
        #     # later we can resume 
        model = pretrain_model.to(device)
        del pretrain_model
    elif POPEN.modual_to_fix is not None:
        # POPEN.model_type != pretrain_popen.model_type
        
        model = POPEN.Model_Class(*POPEN.model_args)
        for modual in POPEN.modual_to_fix:
            if modual in dir(pretrain_model):    
                eval(f'model.{modual}').load_state_dict(
                    eval(f'model.{modual}').state_dict()
                    )

        state_dict = {'epoch': 0,
                        'validation_acc': 0,
                        'state_dict': model.to('cpu'),
                        'validation_loss': np.inf}
        shared_pretrain_pth = POPEN.vae_pth_path.replace(f"_cv{args.kfold_index}", '')
        if not os.path.exists(shared_pretrain_pth):
            utils.snapshot(shared_pretrain_pth, state_dict)
        utils.snapshot(POPEN.vae_pth_path, state_dict)

        model = torch.load(POPEN.vae_pth_path, map_location=torch.device('cpu'))['state_dict']
        model = model.to(device)

    
# -- end2end -- 
elif POPEN.model_type == "CrossStitch_Model":
    backbone = {}
    for t in POPEN.tasks:
        task_popen = Auto_popen(POPEN.backbone_config[t])
        task_model = task_popen.Model_Class(*task_popen.model_args)
        task_model = utils.load_model(task_popen,task_model,logger)
        backbone[t] = task_model.to(device)
    POPEN.model_args = [backbone] + POPEN.model_args
    model = POPEN.Model_Class(*POPEN.model_args).to(device)
else:
    Model_Class = POPEN.Model_Class  # DL_models.LSTM_AE 
    model = Model_Class(*POPEN.model_args).to(device)
    
if POPEN.Resumable:
    model = utils.load_model(POPEN, model, logger)
    
# =========== fix parameters ===========
if isinstance(POPEN.modual_to_fix, list):
    for modual in POPEN.modual_to_fix:
        model = utils.fix_parameter(model,modual)
    logger.info(' \t \t ==============<<< %s part is fixed>>>============== \t \t \n'%POPEN.modual_to_fix)
# =========== set optimizer ===========
if POPEN.optimizer == 'Schedule':
    optimizer = ScheduledOptim(optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
                                      betas=(0.9, 0.98), 
                                      eps=1e-09, 
                                      weight_decay=1e-4, 
                                      amsgrad=True),
                           n_warmup_steps=20)
elif type(POPEN.optimizer) == dict:
    optimizer = eval(scheduleoptim_dict_str.format(**POPEN.optimizer))
else:
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 
                            lr=POPEN.lr,
                            betas=(0.9, 0.98), 
                            eps=1e-09, 
                            weight_decay=POPEN.l2)
if POPEN.loss_schema == 'DTP':
    POPEN.loss_schedualer = Dynamic_Task_Priority(POPEN.tasks,POPEN.gamma,POPEN.chimerla_weight)
elif POPEN.loss_schema == 'DWA':
    POPEN.loss_schedualer = Dynamic_Weight_Averaging(POPEN.tasks,POPEN.tau,POPEN.chimerla_weight)
# =========== resume ===========
best_loss = np.inf
best_acc = 0
best_epoch = 0
previous_epoch = 0
if POPEN.Resumable:
    previous_epoch,best_loss,best_acc = utils.resume(POPEN, optimizer,logger)
    

#                               |=====================================|
#                               |==========  training  part ==========|
#                               |=====================================|
for epoch in range(POPEN.max_epoch-previous_epoch+1):
    epoch += previous_epoch
    
    #           ----------| train |----------
    logger.info("===============================|    epoch {}   |===============================".format(epoch))
    train_val.train(dataloader=train_loader,model=model,optimizer=optimizer,popen=POPEN,epoch=epoch)
           
    #         -----------| validate |-----------
    
    if epoch % POPEN.config_dict['setp_to_check'] == 0:
        logger.info("===============================| start validation |===============================")
        val_total_loss,val_avg_acc = train_val.validate(val_loader,model,popen=POPEN,epoch=epoch)
        _,_ = train_val.validate(test_loader,model,popen=POPEN,epoch=epoch)
        
        DICT ={"ran_epoch":epoch,"n_current_steps":optimizer.n_current_steps,"delta":optimizer.delta} if type(optimizer) == ScheduledOptim else {"ran_epoch":epoch}
        POPEN.update_ini_file(DICT,logger)
        
    #    -----------| compare the result |-----------
        if (best_loss > val_total_loss): #| (best_acc < val_avg_acc):
            # update best performance
            best_loss = min(best_loss,val_total_loss)
            best_acc = max(best_acc,val_avg_acc)
            best_epoch = epoch
            
            # save
            utils.snapshot(POPEN.vae_pth_path, {
                        'epoch': epoch + 1,
                        'validation_acc': val_avg_acc,
                        'state_dict': model.to('cpu'),
                        'validation_loss': val_total_loss,
                        # 'optimizer': optimizer.state_dict(),
                    })
            
            # update the popen
            POPEN.update_ini_file({'run_name':run_name,
                                "ran_epoch":epoch,
                                "best_acc":best_acc,
                                "cuda_id":cuda_id},
                                logger)
            
        elif (epoch - best_epoch >= 30)&((type(optimizer) == ScheduledOptim)):
            optimizer.increase_delta()
            
        elif (epoch - best_epoch >= 60)&(epoch > POPEN.max_epoch/2):
            # at the late phase of training
            logger.info("<<<<<<<<<<< Early Stopping >>>>>>>>>>")
            break