opts.py

import argparse
from config import Constants
import os


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-d',
        '--dataset',
        type=str,
        default='Youtube2Text',
        help='MSRVTT | Youtube2Text',
    )
    parser.add_argument('-m', '--modality', type=str, default='mi')
    parser.add_argument(
        '--cluster', type=int, default=False, help="baseline | baseline+cluster"
    )
    parser.add_argument(
        '--cluster_numbers', type=int, default=5, help='number of clusters'
    )
    parser.add_argument(
        '--base_checkpoint_path',
        type=str,
        default='./experiments_20frame_TRMencoder_SP_C5_SR_0.2',
        help='result folder',
    )
    parser.add_argument('-df', '--default', default=True, action='store_true')
    parser.add_argument('--scope', type=str, default='')
    parser.add_argument('-field', '--field', nargs='+', type=str, default=['seed'])
    parser.add_argument('--no_cuda', default=False, action='store_true')
    parser.add_argument(
        '--method',
        type=str,
        default='ARB',
        help='ARB: autoregressive baseline \n' 'ARB2: ARB + visual word generation',
    )

    parser.add_argument(
        '--encoder',
        type=str,
        default='Transformer',
        help='specify the encoder if we want',
    )  # Encoder_HighWay
    parser.add_argument(
        '--decoder',
        type=str,
        default='BertDecoder',
        help='specify the decoder if we want',
    )
    parser.add_argument(
        '--decoding_type', type=str, default='ARFormer', help='ARFormer | NARFormer'
    )
    parser.add_argument(
        '--fusion',
        type=str,
        default='temporal_concat',
        help='temporal_concat | addition',
    )

    model = parser.add_argument_group(title='Model Parameters')
    # Transformer Configurations
    model.add_argument(
        '--dim_hidden', type=int, default=512, help='size of the rnn hidden layer'
    )
    model.add_argument('--num_hidden_layers_decoder', type=int, default=1)
    model.add_argument('--num_attention_heads', type=int, default=8)
    model.add_argument('--intermediate_size', type=int, default=2048)
    model.add_argument('--hidden_act', type=str, default='gelu_new')
    model.add_argument('--hidden_dropout_prob', type=float, default=0.5)
    model.add_argument('--attention_probs_dropout_prob', type=float, default=0.0)
    model.add_argument("--max_len", type=int, default=30, help='max length of captions')
    model.add_argument('--layer_norm_eps', type=float, default=0.00001)
    model.add_argument('--watch', type=int, default=0)
    model.add_argument('--pos_attention', default=False, action='store_true')
    model.add_argument(
        '--enhance_input',
        type=int,
        default=2,
        help='for NA decoding, 0: without R | 1: re-sampling(R)) | 2: meanpooling(R)',
    )
    model.add_argument('--with_layernorm', default=False, action='store_true')

    model.add_argument(
        '-wc',
        '--with_category',
        default=False,
        action='store_true',
        help='specified for the MSRVTT dataset, use category tags or not',
    )
    model.add_argument('--num_category', type=int, default=20)

    model.add_argument(
        '--encoder_dropout',
        type=float,
        default=0.5,
        help='strength of dropout in the encoder',
    )
    model.add_argument(
        '--no_encoder_bn',
        default=False,
        action='store_true',
        help='by default, a BN layer is placed after the encoder outputs of a modality',
    )
    model.add_argument('--norm_type', type=str, default='bn')
    model.add_argument(
        '--dim_word',
        type=int,
        default=512,
        help='the embedding size of each token in the vocabulary',
    )
    model.add_argument(
        '-tie',
        '--tie_weights',
        default=False,
        action='store_true',
        help='share the weights between word embeddings and the projection layer',
    )

    training = parser.add_argument_group(title='Training Parameters')
    training.add_argument('--seed', default=0, type=int, help='for reproducibility')
    training.add_argument(
        '--learning_rate', default=5e-4, type=float, help='the initial larning rate'
    )
    training.add_argument(
        '--decay',
        default=0.9,
        type=float,
        help='the decay rate of learning rate per epoch',
    )
    training.add_argument(
        '--minimum_learning_rate',
        default=5e-5,
        type=float,
        help='the minimum learning rate',
    )
    training.add_argument(
        '--n_warmup_steps',
        type=int,
        default=0,
        help='the number of warmup steps towards the initial lr',
    )
    training.add_argument('--optim', type=str, default='adam', help='adam | rmsprop')
    training.add_argument(
        '--grad_clip', type=float, default=5, help='clip gradients at this value'
    )
    training.add_argument(
        '--weight_decay',
        type=float,
        default=5e-4,
        help='Strength of weight regularization',
    )
    training.add_argument(
        '-e', '--epochs', type=int, default=20, help='number of epochs'
    )
    training.add_argument(
        '-b', '--batch_size', type=int, default=64, help='minibatch size'
    )

    training.add_argument(
        '--pretrained_path', default='', type=str, help='path for the pretrained model'
    )
    # NA decoding
    training.add_argument(
        '--teacher_path', type=str, default='', help='path for the AR-B model'
    )
    training.add_argument(
        '- -beta',
        nargs='+',
        type=float,
        default=[0, 1],
        help='len=2, [lowest masking ratio, highest masking ratio]',
    )
    training.add_argument(
        '--visual_word_generation', default=False, action='store_true'
    )
    training.add_argument(
        '--demand',
        nargs='+',
        type=str,
        default=['VERB', 'NOUN'],
        help='pos_tag we want to focus on when training with visual word generation',
    )
    training.add_argument(
        '-nvw',
        '--nv_weights',
        nargs='+',
        type=float,
        default=[0.8, 1.0],
        help='len=2, weights of visual word generation and caption generation (or mlm)',
    )
    training.add_argument(
        '--load_teacher_weights',
        default=False,
        action='store_true',
        help='specified for NA-based models, initialize randomly or inherit from the teacher (AR-B)',
    )
    training.add_argument('--no_test', default=False, action='store_true')

    evaluation = parser.add_argument_group(title='Evaluation Parameters')
    evaluation.add_argument(
        '-see',
        '--start_eval_epoch',
        type=int,
        default=0,
        help='start evaluation after a specific epoch',
    )
    evaluation.add_argument(
        '--tolerence', type=int, default=1000, help='for early stop'
    )
    evaluation.add_argument(
        '--metric_sum',
        nargs='+',
        type=int,
        default=[1, 1, 1, 1],
        help='meta sum of the metrics',
    )
    evaluation.add_argument(
        '--standard',
        nargs='+',
        type=str,
        default=[
            'Bleu_4',
            'METEOR',
            'CIDEr',
        ],  # ['Bleu_4', 'METEOR', 'ROUGE_L', 'CIDEr'],
        help='the standard of performance to select the best model',
    )
    evaluation.add_argument(
        '-bs',
        '--beam_size',
        type=int,
        default=1,
        help='specified for AR decoding, the number of candidates',
    )
    evaluation.add_argument(
        '-ba',
        '--beam_alpha',
        type=float,
        default=1.0,
        help='the preference of sentence length, larger --> longer',
    )
    # NA decoding
    evaluation.add_argument(
        '--paradigm',
        type=str,
        default='mp',
        help='mp: MaskPredict | l2r: Left2Right | ef: EasyFirst',
    )
    evaluation.add_argument(
        '-lbs',
        '--length_beam_size',
        type=int,
        default=6,
        help='specified for NA decoding, the number of length candidates',
    )
    evaluation.add_argument(
        '--iterations',
        type=int,
        default=5,
        help='the number of iterations for the MP algorithm',
    )
    evaluation.add_argument(
        '--q',
        type=int,
        default=1,
        help='the number of tokens to update for L2R & EF algorithms',
    )
    evaluation.add_argument(
        '--q_iterations',
        type=int,
        default=1,
        help='the number of iterations for L2R & EF algorithms',
    )
    evaluation.add_argument(
        '--use_ct',
        default=False,
        action='store_true',
        help='use coarse-grained templates or not, only for methods with visual word generation',
    )
    # checkpoint settings
    evaluation.add_argument(
        '--k_best_model',
        type=int,
        default=1,
        help='checkpoints with top-k performance will be saved',
    )
    evaluation.add_argument(
        '--save_checkpoint_every',
        type=int,
        default=1,
        help='how often to save a model checkpoint (in epoch)?',
    )

    multitask = parser.add_argument_group(title='Multi-Task Parameters')
    multitask.add_argument(
        '--crit', nargs='+', type=str, default=['lang'], help='lang | length'
    )
    multitask.add_argument('--crit_name', nargs='+', type=str, default=['Cap Loss'])
    multitask.add_argument('--crit_scale', nargs='+', type=float, default=[1.0])

    dataloader = parser.add_argument_group(title='Dataloader Parameters')
    dataloader.add_argument(
        '--n_frames',
        type=int,
        default=20,
        help='the number of frames to represent a whole video',
    )
    dataloader.add_argument(
        '--n_caps_per_video',
        type=int,
        default=0,
        help='the number of captions per video to constitute the training set',
    )
    dataloader.add_argument(
        '--random_type',
        type=str,
        default='segment_random',
        help='sampling strategy during training: segment_random (default) | all_random | equally_sampling',
    )
    dataloader.add_argument(
        '--load_feats_type',
        type=int,
        default=1,
        help='load feats from the same frame_ids (0) '
        'or different frame_ids (1), '
        'or directly load all feats without sampling (2)',
    )

    # modality information
    dataloader.add_argument(
        '--dim_a', type=int, default=1, help='feature dimension of the audio modality'
    )
    dataloader.add_argument(
        '--dim_m',
        type=int,
        default=2048,
        help='feature dimension of the motion modality',
    )
    dataloader.add_argument(
        '--dim_i',
        type=int,
        default=2048,
        help='feature dimension of the image modality',
    )
    dataloader.add_argument(
        '--dim_o', type=int, default=1, help='feature dimension of the object modality'
    )
    dataloader.add_argument('--dim_t', type=int, default=1)
    dataloader.add_argument('--feats_a_name', nargs='+', type=str, default=[])
    dataloader.add_argument(
        '--feats_m_name',
        nargs='+',
        type=str,
        default=['motion_resnext101_kinetics_duration16_overlap8.hdf5'],
    )
    dataloader.add_argument(
        '--feats_i_name',
        nargs='+',
        type=str,
        default=['image_resnet101_imagenet_fps_max60.hdf5'],
    )
    dataloader.add_argument('--feats_o_name', nargs='+', type=str, default=[])
    dataloader.add_argument('--feats_t_name', nargs='+', type=str, default=[])
    # corpus information
    dataloader.add_argument('--info_corpus_name', type=str, default='info_corpus.pkl')
    dataloader.add_argument('--reference_name', type=str, default='refs.pkl')

    args = parser.parse_args()
    check_dataset(args)
    check_method(args)
    check_valid(args)
    return args


def check_valid(args):
    assert args.load_feats_type in [0, 1, 2]
    if not args.default:
        assert (
            args.scope
        ), "Please add the argument \'--scope $folder_name_to_save_models\'"


def check_dataset(args):
    if args.dataset.lower() == 'msvd':
        args.dataset = 'Youtube2Text'

    assert args.dataset in [
        'Youtube2Text',
        'MSRVTT',
    ], "We now only support Youtube2Text (MSVD) and MSRVTT datasets."

    if args.default:
        if args.dataset == 'Youtube2Text':
            args.beta = [0, 1]
            args.max_len = 20
            args.with_category = False
        elif args.dataset == 'MSRVTT':
            args.beta = [0.35, 0.9]
            args.max_len = 30
            args.with_category = True

    if args.dataset == 'Youtube2Text':
        assert (
            not args.with_category
        ), "Category information is not available in the Youtube2Text (MSVD) dataset"


def check_method(args):
    if args.method:
        import yaml

        methods = yaml.full_load(open('./config/methods.yaml'))
        assert (
            args.method in methods.keys()
        ), "The method {} can not be found in ./config/methods.yaml".format(args.method)
        for k, v in methods[args.method].items():
            setattr(args, k, v)

    if args.decoding_type == 'NARFormer':
        args.crit = ['lang', 'length']
        args.crit_name = ['Cap Loss', 'Length Loss']
        args.crit_scale = [1.0, 1.0]
    args.crit_key = [Constants.mapping[item.lower()] for item in args.crit]

    if args.default:
        if args.decoding_type == 'NARFormer':
            if args.visual_word_generation:
                args.use_ct = True
                args.nv_weights = [0.8, 1.0]
            args.enhance_input = 2
            args.length_beam_size = int(6)
            args.iterations = int(5)
            args.beam_alpha = 1.35 if args.dataset == 'MSRVTT' else 1.0
            args.algorithm_print_sent = True
            args.teacher_path = os.path.join(
                args.base_checkpoint_path,
                args.dataset,
                'ARB',
                args.scope,
                'best.pth.tar',
            )
            assert os.path.exists(args.teacher_path)
            args.load_teacher_weights = True
            args.with_teacher = True
        else:
            args.beam_size = int(5.0)
            args.beam_alpha = 1.0