main/config.py

import os
import pdb
import time
import torch
import logging
import argparse
import importlib
from utils.basic_utils import mkdirp, remkdirp, \
    load_json, save_json, make_zipfile, dict_to_markdown

logger = logging.getLogger(__name__)
logging.basicConfig(format="%(asctime)s.%(msecs)03d:%(levelname)s:%(name)s - %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO)

class BaseOptions(object):
    saved_option_filename = "opt.json"
    ckpt_filename = "model.ckpt"
    tensorboard_log_dir = "tensorboard_log"
    train_log_filename = "train.log.txt"
    eval_log_filename = "eval.log.txt"

    def __init__(self):
        self.parser = None
        self.initialized = False
        self.opt = None

    def initialize(self):
        self.initialized = True
        parser = argparse.ArgumentParser()
        # * Running configs
        parser.add_argument("--dset_type", type=str, choices=["mr", "hl", "vs", "vlp"])    # moment retrieval, highlight detection, and video summarization
        parser.add_argument("--dset_name", type=str, choices=["qvhighlights", "charades", "anet", "tvsum", "youtube", "summe", "ego4d", "qfvs", "video2gif", "coin", "hacs", "vlp", "videocc", "tacos"])
        parser.add_argument("--domain_name", type=str, default=None)
        parser.add_argument("--model_id", type=str, default="moment_detr")
        parser.add_argument("--exp_id", type=str, default="debug", help="id of this run, required at training")
        parser.add_argument("--device", type=int, default=0, help="0 cuda, -1 cpu")
        parser.add_argument("--gpu_id", type=int, default=0)
        parser.add_argument("--debug", action="store_true",
                            help="debug (fast) mode, break all loops, do not load all data into memory.")
        parser.add_argument("--seed", type=int, default=2018, help="random seed")

        # * DDP
        parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')


        parser.add_argument("--eval_split_name", type=str, default="val",
                            help="should match keys in video_duration_idx_path, must set for VCMR")
        parser.add_argument("--data_ratio", type=float, default=1.0,
                            help="how many training and eval data to use. 1.0: use all, 0.1: use 10%."
                                 "Use small portion for debug purposes. Note this is different from --debug, "
                                 "which works by breaking the loops, typically they are not used together.")
        parser.add_argument("--results_root", type=str, default="results")
        parser.add_argument("--num_workers", type=int, default=0,
                            help="num subprocesses used to load the data, 0: use main process")
        parser.add_argument("--no_pin_memory", action="store_true",
                            help="Don't use pin_memory=True for dataloader. "
                                 "ref: https://discuss.pytorch.org/t/should-we-set-non-blocking-to-true/38234/4")

        # * Training configs
        parser.add_argument("--bsz", type=int, default=32, help="mini-batch size")
        parser.add_argument("--n_epoch", type=int, default=200, help="number of epochs to run")
        parser.add_argument("--max_es_cnt", type=int, default=200,
                            help="number of epochs to early stop, use -1 to disable early stop")
        parser.add_argument("--lr", type=float, default=1e-4, help="learning rate")
        parser.add_argument("--lr_drop", type=int, default=400, help="drop learning rate to 1/10 every lr_drop epochs")
        parser.add_argument("--lr_gamma", type=float, default=0.1, help="lr reduces the gamma times after the `drop' epoch")
        parser.add_argument("--lr_warmup", type=float, default=-1, help="linear warmup scheme")
        parser.add_argument("--wd", type=float, default=1e-4, help="weight decay")
        parser.add_argument("--grad_clip", type=float, default=0.1, help="perform gradient clip, -1: disable")

        # ** Loss coefficients
        # *** boundary branch
        parser.add_argument("--span_loss_type", default="l1", type=str, choices=['l1', 'ce'],
                            help="l1: (center-x, width) regression. ce: (st_idx, ed_idx) classification.")
        parser.add_argument('--b_loss_coef', default=10, type=float)    # boundary regression e.g., l1
        parser.add_argument('--g_loss_coef', default=1, type=float) # giou loss
        # *** foreground branch
        parser.add_argument('--eos_coef', default=0.1, type=float, help="relative classification weight of the no-object class")
        parser.add_argument('--f_loss_coef', default=4, type=float) # cls loss for foreground
        # *** saliency branch
        parser.add_argument("--s_loss_intra_coef", type=float, default=1., help="inter-video (frame-level) saliency loss e.g. momentdetr saliency loss")
        parser.add_argument("--s_loss_inter_coef", type=float, default=0., help="intra-video (sample-level) saliency loss,")

        # * Eval configs
        parser.add_argument("--main_metric", type=str, default="MR-full-mAP")
        parser.add_argument('--eval_mode', default=None, type=str,
                            help="how to integrate foreground and saliency for better prediction")
        parser.add_argument("--eval_bsz", type=int, default=100,
                            help="mini-batch size at inference, for query")
        parser.add_argument("--eval_epoch", type=int, default=5,
                            help="number of epochs for once inference")
        parser.add_argument("--eval_init", action="store_true", help="evaluate model before training i.e. `epoch=-1'")
        parser.add_argument("--save_interval", type=int, default=50)

        parser.add_argument("--resume", type=str, default=None,
                            help="checkpoint path to resume or evaluate, without --resume_all this only load weights")
        parser.add_argument("--resume_dir", type=str, default=None,
                            help="checkpoint path to resume or evaluate, without --resume_all this only load weights")
        parser.add_argument("--resume_all", action="store_true",
                            help="if --resume_all, load optimizer/scheduler/epoch as well")
        parser.add_argument("--start_epoch", type=int, default=None,
                            help="if None, will be set automatically when using --resume_all")

        # ** NMS configs
        parser.add_argument("--no_sort_results", action="store_true",
                            help="do not sort results, use this for moment query visualization")
        parser.add_argument("--max_before_nms", type=int, default=10)
        parser.add_argument("--max_after_nms", type=int, default=10)
        parser.add_argument("--conf_thd", type=float, default=0.0, help="only keep windows with conf >= conf_thd")
        parser.add_argument("--nms_thd", type=float, default=-1,
                            help="additionally use non-maximum suppression "
                                 "(or non-minimum suppression for distance)"
                                 "to post-processing the predictions. "
                                 "-1: do not use nms. [0, 1]")

        # * Dataset configs
        parser.add_argument("--use_cache",  type=int, default=-1, help="Preload features into cache for fast IO")
        parser.add_argument("--max_q_l", type=int, default=75)
        parser.add_argument("--max_v_l", type=int, default=75)
        parser.add_argument("--clip_length", type=float, default=1.0)
        parser.add_argument("--clip_len_list", type=int, nargs='+')
        parser.add_argument("--max_windows", type=int, default=5)

        parser.add_argument("--add_easy_negative", type=int, default=1)
        parser.add_argument("--easy_negative_only", type=int, default=1)
        parser.add_argument("--round_multiple", type=int, default=1)

        parser.add_argument("--train_path", type=str, default=None, nargs='+')
        parser.add_argument("--eval_path", type=str, default=None,
                            help="Evaluating during training, for Dev set. If None, will only do training, ")
        parser.add_argument("--train_path_list", type=str, nargs='+')
        parser.add_argument("--eval_path_list", type=str, nargs='+')
        parser.add_argument("--feat_root_list", type=str, nargs='+')

        parser.add_argument("--no_norm_vfeat", action="store_true", help="Do not do normalize video feat")
        parser.add_argument("--no_norm_tfeat", action="store_true", help="Do not do normalize text feat")
        parser.add_argument("--v_feat_dirs", type=str, nargs="+",
                            help="video feature dirs. If more than one, will concat their features. "
                                 "Note that sub ctx features are also accepted here.")
        parser.add_argument("--t_feat_dir", type=str, help="text/query feature dir")
        parser.add_argument("--v_feat_dim", type=int, help="video feature dim")
        parser.add_argument("--t_feat_dim", type=int, help="text/query feature dim")
        parser.add_argument("--ctx_mode", type=str, default="video_tef")
        parser.add_argument("--v_feat_types", type=str)
        parser.add_argument("--t_feat_type", type=str)

        # * Model configs
        parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),
                            help="Type of positional embedding to use on top of the image features")
        parser.add_argument("--n_input_proj", type=int, default=2, help="#layers to vid/txt projector")
        parser.add_argument("--temperature", type=float, default=0.07, help="temperature nce contrastive_align_loss")

        # ** Transformer
        parser.add_argument('--enc_layers', default=4, type=int,
                            help="Number of encoding layers in the transformer")
        parser.add_argument('--sub_enc_layers', default=2, type=int,
                            help="Number of encoding layers in the video / text transformer in albef-style.")
        parser.add_argument('--dec_layers', default=2, type=int,
                            help="Number of decoding layers in the transformer, N/A for UniVTG")
        parser.add_argument('--dim_feedforward', default=1024, type=int,
                            help="Intermediate size of the feedforward layers in the transformer blocks")
        parser.add_argument('--hidden_dim', default=256, type=int,
                            help="Size of the embeddings (dimension of the transformer)")
        parser.add_argument('--input_dropout', default=0.5, type=float,
                            help="Dropout applied in input")
        parser.add_argument('--dropout', default=0.1, type=float,
                            help="Dropout applied in the transformer")
        parser.add_argument('--droppath', default=0.1, type=float,
                            help="Droppath applied in the transformer")
        parser.add_argument("--txt_drop_ratio", default=0, type=float,
                            help="drop txt_drop_ratio tokens from text input. 0.1=10%")
        parser.add_argument("--use_txt_pos", action="store_true", help="use position_embedding for text as well.")
        parser.add_argument('--nheads', default=8, type=int,
                            help="Number of attention heads inside the transformer's attentions")
        parser.add_argument('--num_queries', default=10, type=int,
                            help="Number of query slots")
        parser.add_argument('--pre_norm', action='store_true')

        # ** momentdetr configs e.g. Matcher, saliency margin
        parser.add_argument('--set_cost_span', default=10, type=float,
                            help="L1 span coefficient in the matching cost")
        parser.add_argument('--set_cost_giou', default=1, type=float,
                            help="giou span coefficient in the matching cost")
        parser.add_argument('--set_cost_class', default=4, type=float,
                            help="Class coefficient in the matching cost")
        parser.add_argument("--saliency_margin", type=float, default=0.2)
        parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_true',
                            help="Disables auxiliary decoding losses (loss at each layer)")

        # * Query-Force Video Summarization
        parser.add_argument("--max_segment_num", type=int, default=20)
        parser.add_argument("--max_frame_num", type=int, default=200)
        parser.add_argument("--top_percent", type=float, default=0.02)

        parser.add_argument("--qfvs_vid_feature", type=str, default='fps1')
        parser.add_argument("--qfvs_txt_feature", type=str, default='query')
        parser.add_argument("--qfvs_split", type=int, default=-1)

        parser.add_argument("--qfvs_dense_shot", type=int, default=-1)
        parser.add_argument("--qfvs_score_ensemble", type=int, default=-1)
        parser.add_argument("--qfvs_score_gather", type=int, default=-1)
        parser.add_argument("--qfvs_loss_gather", type=int, default=-1)
        self.parser = parser

    def display_save(self, opt):
        args = vars(opt)
        # Display settings
        print(dict_to_markdown(vars(opt), max_str_len=120))
        # Save settings
        if not isinstance(self, TestOptions):
            option_file_path = os.path.join(opt.results_dir, self.saved_option_filename)  # not yaml file indeed
            save_json(args, option_file_path, save_pretty=True)

    def parse(self, args=None):
        if not self.initialized:
            self.initialize()
        opt = self.parser.parse_args()
        
        if args is not None:
            args_dict = vars(args)
            opt_dict = vars(opt)
            for key, value in args_dict.items():
                opt_dict[key] = value
            opt = argparse.Namespace(**opt_dict)    
            opt.model_dir = os.path.dirname(opt.resume)
            torch.cuda.set_device(opt.gpu_id)
            
        if opt.debug:
            opt.results_root = os.path.sep.join(opt.results_root.split(os.path.sep)[:-1] + ["debug_results", ])
            opt.num_workers = 0

        if isinstance(self, TestOptions):
            # modify model_dir to absolute path
            # opt.model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results", opt.model_dir)
            opt.model_dir = os.path.dirname(opt.resume)
            saved_options = load_json(os.path.join(opt.model_dir, self.saved_option_filename))
            for arg in saved_options:  # use saved options to overwrite all BaseOptions args.
                if arg not in ["results_root", "num_workers", "nms_thd", "debug",  "max_before_nms", "max_after_nms"
                               "max_pred_l", "min_pred_l", "gpu_id",
                               "resume", "resume_all", "no_sort_results",
                               "eval_path", "eval_split_name"]:
                            #    "dset_name", "v_feat_dirs", "t_feat_dir"]:
                    setattr(opt, arg, saved_options[arg])
            # opt.no_core_driver = True
            if opt.eval_results_dir is not None:
                opt.results_dir = opt.eval_results_dir
        else:
            if opt.exp_id is None:
                raise ValueError("--exp_id is required for at a training option!")

            # ctx_str = opt.ctx_mode + "_sub" if any(["sub_ctx" in p for p in opt.v_feat_dirs]) else opt.ctx_mode

            if 'debug' not in opt.exp_id:
                opt.results_dir = os.path.join(opt.results_root, "-".join([opt.dset_type, opt.dset_name]), "-".join([opt.exp_id, opt.v_feat_types, opt.t_feat_type, time.strftime("%Y_%m_%d_%H")]))
            else:
                opt.results_dir = os.path.join(opt.results_root, "-".join([opt.dset_type, opt.dset_name]), opt.exp_id) # debug mode.

            if int(opt.local_rank) in [0, -1]:
                # mkdirp(opt.results_dir)
                remkdirp(opt.results_dir)   # remove dir and remkdir it.

                # save a copy of current code
                code_dir = os.path.dirname(os.path.realpath(__file__))
                code_zip_filename = os.path.join(opt.results_dir, "code.zip")
                make_zipfile(code_dir, code_zip_filename,
                            enclosing_dir="code",
                            exclude_dirs_substring="results",
                            exclude_dirs=["results", "debug_results", "__pycache__"],
                            exclude_extensions=[".pyc", ".ipynb", ".swap"], )

        if int(opt.local_rank) in [0, -1]:
            self.display_save(opt)
            opt.ckpt_filepath = os.path.join(opt.results_dir, self.ckpt_filename)
            opt.train_log_filepath = os.path.join(opt.results_dir, self.train_log_filename)
            opt.eval_log_filepath = os.path.join(opt.results_dir, self.eval_log_filename)
            opt.tensorboard_log_dir = os.path.join(opt.results_dir, self.tensorboard_log_dir)
            # opt.device = torch.device("cuda" if opt.device >= 0 else "cpu")

        if int(opt.local_rank) in [-1]:
            torch.cuda.set_device(opt.gpu_id)
        opt.pin_memory = not opt.no_pin_memory

        if opt.local_rank == -1:
            torch.cuda.set_device(opt.gpu_id)

        opt.use_tef = "tef" in opt.ctx_mode
        opt.use_video = "video" in opt.ctx_mode
        if not opt.use_video:
            opt.v_feat_dim = 0
        if opt.use_tef:
            opt.v_feat_dim += 2

        self.opt = opt
        return opt

class TestOptions(BaseOptions):
    """add additional options for evaluating"""

    def initialize(self):
        BaseOptions.initialize(self)
        # also need to specify --eval_split_name
        self.parser.add_argument("--eval_id", type=str, help="evaluation id")
        self.parser.add_argument("--eval_results_dir", type=str, default=None,
                                 help="dir to save results, if not set, fall back to training results_dir")
        self.parser.add_argument("--model_dir", type=str,
                                 help="dir contains the model file, will be converted to absolute path afterwards")

class WarmupStepLR(torch.optim.lr_scheduler.StepLR):
    def __init__(self, optimizer, warmup_steps, step_size, gamma=0.1, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.step_size = step_size
        self.gamma = gamma
        super(WarmupStepLR, self).__init__(optimizer, step_size, gamma=self.gamma, last_epoch=last_epoch)
    def get_lr(self):
        if not self._get_lr_called_within_step:
            import warnings
            warnings.warn("To get the last learning rate computed by the scheduler, "
                          "please use `get_last_lr()`.", DeprecationWarning)
        # e.g. warmup_steps = 10, case: 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 21...
        if self.last_epoch == self.warmup_steps or(self.last_epoch % self.step_size != 0 and self.last_epoch > self.warmup_steps):
            return [group['lr'] for group in self.optimizer.param_groups]
        # e.g. warmup_steps = 10, case: 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
        elif self.last_epoch < self.warmup_steps:
            return [group['initial_lr'] * float(self.last_epoch + 1) / float(self.warmup_steps) for group in self.optimizer.param_groups]
        
        
        # e.g. warmup_steps = 10, case: 10, 20, 30, 40...
        return [group['lr'] * self.gamma
                for group in self.optimizer.param_groups]
    def _get_closed_form_lr(self):
        if self.last_epoch <= self.warmup_steps:
            return [base_lr * float(self.last_epoch) / (self.warmup_steps) for base_lr in self.base_lrs]
        else:
            return [base_lr * self.gamma ** ((self.last_epoch -  self.warmup_steps)// self.step_size) for base_lr in self.base_lrs]

def setup_model(opt):
    """setup model/optimizer/scheduler and load checkpoints when needed"""
    logger.info("setup model/optimizer/scheduler")

    importer = importlib.import_module('.'.join(['model', opt.model_id]))
    model, criterion = importer.build_model(opt)

    if int(opt.device) >= 0:
        logger.info("CUDA enabled.")
        model.to(opt.gpu_id)
        criterion.to(opt.gpu_id)

    param_dicts = [{"params": [p for n, p in model.named_parameters() if p.requires_grad]}]
    optimizer = torch.optim.AdamW(param_dicts, lr=opt.lr, weight_decay=opt.wd)

    if opt.lr_warmup != -1 and opt.lr_drop > 0:
        lr_scheduler = WarmupStepLR(optimizer, warmup_steps=opt.lr_warmup[0], step_size=opt.lr_drop, gamma=opt.lr_gamma)
    
    elif opt.lr_warmup != -1:
        from transformers import get_constant_schedule_with_warmup
        lr_scheduler =  get_constant_schedule_with_warmup(optimizer, opt.lr_warmup[0])

    elif opt.lr_drop > 0:
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, opt.lr_drop, gamma=opt.lr_gamma)

    if opt.resume is not None:
        logger.info(f"Load checkpoint from {opt.resume}")
        checkpoint = torch.load(opt.resume, map_location="cpu")
        
        for key in list(checkpoint["model"].keys()):
            checkpoint["model"][key.replace('module.', '')] = checkpoint["model"].pop(key)
        model.load_state_dict(checkpoint["model"])

        if opt.resume_all:
            optimizer.load_state_dict(checkpoint['optimizer'])
            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
            opt.start_epoch = checkpoint['epoch'] + 1
        logger.info(f"Loaded model saved at epoch {checkpoint['epoch']} from checkpoint: {opt.resume}")
    else:
        logger.warning("If you intend to evaluate the model, please specify --resume with ckpt path")

    return model, criterion, optimizer, lr_scheduler