run.py

'''
 * Based on coda prompt here
 * https://github.com/GT-RIPL/CODA-Prompt
 * Build our CDL model on CODAPrompt baseline(DualPrompt and L2P)
'''

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import os
import sys
import argparse
import torch
import numpy as np
import yaml
import json
import random
from trainer import Trainer

def create_args():
    
    parser = argparse.ArgumentParser()

    # Standard Args
    parser.add_argument('--random_s', type=int, default=1, help="The random seed")
    parser.add_argument('--gpuid', nargs="+", type=int, default=[0],
                         help="The list of gpuid, ex:--gpuid 3 1. Negative value means cpu-only")
    parser.add_argument('--log_dir', type=str, default="outputs/out",
                         help="Save experiments results in dir for future plotting!")
    parser.add_argument('--learner_type', type=str, default='default', help="The type (filename) of learner")
    parser.add_argument('--learner_name', type=str, default='NormalNN', help="The class name of learner")
    parser.add_argument('--debug_mode', type=int, default=0, metavar='N',
                        help="activate learner specific settings for debug_mode")

    parser.add_argument('--overwrite', type=int, default=0, metavar='N', help='Train regardless of whether saved model exists')

    # CL Args          
    parser.add_argument('--upper_bound_flag', default=False, action='store_true', help='Upper bound')
    parser.add_argument('--memory', type=int, default=0, help="size of memory for replay")
    parser.add_argument('--DW', default=False, action='store_true', help='dataset balancing')
    parser.add_argument('--prompt_param', nargs="+", type=float, default=[1, 1, 1],
                         help="e prompt pool size, e prompt length, g prompt length")
    
    # The Teacher and Student Model
    
    parser.add_argument('--t_model', default='vit_base_patch16_224', type=str, metavar='MODEL', help='Name of t_model to train')
    parser.add_argument('--s_model', default='vit_tiny_patch16_224', type=str, metavar='MODEL', help='Name of s_model to train')
    
    # KD Args
    parser.add_argument('--kd_alpha', type=float, default=0.5, help="alpha of distillation loss")
    parser.add_argument('--Soft_T', type=float, default=2., help="temperature for distillation")
    

    # Config Arg
    parser.add_argument('--config', type=str, default="configs/config.yaml",
                         help="yaml experiment config input")

    # Config the KD methods
    parser.add_argument('--KD_method', type=str, default='KD', help="The KD methods")

    return parser

def get_args(argv):
    parser=create_args()
    args = parser.parse_args(argv)
    config = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)
    config.update(vars(args))
    return argparse.Namespace(**config)

# want to save everything printed to outfile
class Logger(object):
    def __init__(self, name):
        self.terminal = sys.stdout
        self.log = open(name, "a")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)  

    def flush(self):
        self.log.flush()

if __name__ == '__main__':
    args = get_args(sys.argv[1:])

    # determinstic backend
    torch.backends.cudnn.deterministic=True

    # duplicate output stream to output file
    if not os.path.exists(args.log_dir): os.makedirs(args.log_dir)
    log_out = args.log_dir + '/output.log'
    sys.stdout = Logger(log_out)

    # save args
    with open(args.log_dir + '/args.yaml', 'w') as yaml_file:
        yaml.dump(vars(args), yaml_file, default_flow_style=False)
    
    metric_keys = ['acc','time',]



    print('************************************')
    print('* START TRAINING ')
    print('************************************')

    # set random seeds
    seed = args.random_s
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    # set up a trainer
    trainer = Trainer(args, seed, metric_keys)


    # init total run metrics storage
    max_task = trainer.max_task

    # train model
    trainer.train(args)