text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import unicode_literals
  from __future__ import print_function
  
  from torch.utils.data import (SequentialSampler)
  
  from dataloaders.dataloader_meta_msrvtt import MSRVTT_1KA_Test, MSRVTT_7k_Test, MSRVTT_MIX_Train, MSRVTT_7k_Train, MSRVTT_1KA_Train
  from dataloaders.dataloader_meta_youcook import Youcook_Train, Youcook_Retrieval_Test, Youcook_Caption_Test
  
  
  from modules.metrics import compute_metrics
  from nlgeval import NLGEval
  from tqdm import tqdm
  import numpy as np
  import argparse
  import random
  import torch
  import time
  import os
  
  
  from modules.modeling_auxiLearn import MonoLinearHyperNet, MonoNonlinearHyperNet, MTN_auxi
  from modules.optimization_auxiLearn import MetaOptimizer
  from modules.tokenization import BertTokenizer
  from modules.file_utils import PYTORCH_PRETRAINED_BERT_CACHE
  from modules.modeling_MTN import UniVL
  from modules.optimization_meta import BertAdam
  from torch.utils.data import DataLoader
  
  from modules.utils import str2list, AverageMeter, str2bool, parallel_apply, get_logger
  
  
  import wandb
  
  from mtn_caption import eval_caption_epoch
  
  global logger
  import socket
  import torch.multiprocessing as mp
  
  
  
  def get_args(description='UniVL on Retrieval Task'):
      parser = argparse.ArgumentParser(description=description)
      parser.add_argument("--do_pretrain", action='store_true', help="Whether to run training.")
      parser.add_argument("--do_train", action='store_true', help="Whether to run training.")
      parser.add_argument("--do_eval", action='store_true', help="Whether to run eval on the dev set.")
  
      parser.add_argument('--youcook_train_csv', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_train.csv', help='')
      parser.add_argument('--youcook_val_csv', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_val.csv', help='')
      parser.add_argument('--youcook_data_path', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_data.no_transcript.pickle', help='data pickle file path')
      parser.add_argument('--youcook_features_path', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_videos_features.pickle', help='feature path')
  
  
      parser.add_argument('--VTT_train_csv', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_train.9k.csv', help='')
      parser.add_argument('--VTT_val_csv', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_JSFUSION_test.csv', help='')
      # parser.add_argument('--VTT_data_path', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_data.json', help='data pickle file path')
      parser.add_argument('--VTT_data_path', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_transcript_data.json', help='data pickle file path')
      parser.add_argument('--VTT_features_path', type=str, default='../../../../../data/project/rw/msrvtt_feature/msrvtt_videos_features.pickle', help='feature path')
      parser.add_argument('--VTT_train_1KB_csv', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_train.1KB.csv', help='')
  
  
  
      parser.add_argument("--init_model", default="/data/project/rw/joonmyung/data/UniVL/univl.pretrained.bin", type=str, required=False, help="Initial model.")
      parser.add_argument('--num_thread_reader', type=int, default=8, help='')
      parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate')
      parser.add_argument('--epochs', type=int, default=20, help='upper epoch limit')
      parser.add_argument('--batch_size', type=int, default=128, help='batch size')
      parser.add_argument('--batch_size_val', type=int, default=64, help='batch size eval')
      parser.add_argument('--lr_decay', type=float, default=0.9, help='Learning rate exp epoch decay')
      parser.add_argument('--n_display', type=int, default=10, help='Information display frequence')
      parser.add_argument('--video_dim', type=int, default=1024, help='video feature dimension')
      parser.add_argument('--seed', type=int, default=42, help='random seed')
      parser.add_argument('--max_words', type=int, default=20, help='')
      parser.add_argument('--max_frames', type=int, default=100, help='')
      parser.add_argument('--feature_framerate', type=int, default=1, help='')
      parser.add_argument('--margin', type=float, default=0.1, help='margin for loss')
      parser.add_argument('--hard_negative_rate', type=float, default=0.5, help='rate of intra negative sample')
      parser.add_argument('--negative_weighting', type=int, default=1, help='Weight the loss for intra negative')
      parser.add_argument('--n_pair', type=int, default=1, help='Num of pair to output from data loader')
  
      parser.add_argument("--output_dir", default=None, type=str, required=True,
                          help="The output directory where the model predictions and checkpoints will be written.")
      parser.add_argument("--bert_model", default="bert-base-uncased", type=str, required=True,
                          help="Bert pre-trained model")
      parser.add_argument("--visual_model", default="visual-base", type=str, required=False, help="Visual module")
      parser.add_argument("--cross_model", default="cross-base", type=str, required=False, help="Cross module")
      parser.add_argument("--decoder_model", default="decoder-base", type=str, required=False, help="Decoder module")
      parser.add_argument("--do_lower_case", action='store_true', help="Set this flag if you are using an uncased model.")
      parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                          help="Number of updates steps to accumulate before performing a backward/update pass.")
      parser.add_argument('--n_gpu', type=int, default=1, help="Changed in the execute process.")
  
      parser.add_argument("--cache_dir", default="", type=str,
                          help="Where do you want to store the pre-trained models downloaded from s3")
  
      parser.add_argument('--fp16', action='store_true',
                          help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
      parser.add_argument('--fp16_opt_level', type=str, default='O1',
                          help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                               "See details at https://nvidia.github.io/apex/amp.html")
  
      parser.add_argument("--task_type", default="retrieval", type=str, help="Point the task `retrieval` to finetune.")
  
  
      parser.add_argument("--world_size", default=0, type=int, help="distribted training")
      parser.add_argument("--local_rank", default=0, type=int, help="distribted training")
      parser.add_argument('--coef_lr', type=float, default=0.1, help='coefficient for bert branch.')
      parser.add_argument('--use_mil', action='store_true', help="Whether use MIL as Miech et. al. (2020).")
      parser.add_argument('--sampled_use_mil', action='store_true', help="Whether MIL, has a high priority than use_mil.")
  
      parser.add_argument('--text_num_hidden_layers', type=int, default=12, help="Layer NO. of text.")
      parser.add_argument('--visual_num_hidden_layers', type=int, default=6, help="Layer NO. of visual.")
      parser.add_argument('--cross_num_hidden_layers', type=int, default=2, help="Layer NO. of cross.")
      parser.add_argument('--decoder_num_hidden_layers', type=int, default=3, help="Layer NO. of decoder.")
  
      parser.add_argument('--train_sim_after_cross', action='store_true', help="Test retrieval after cross encoder.")
      parser.add_argument('--expand_msrvtt_sentences', action='store_true', help="")
  
  
  
      parser.add_argument('--use_wandb', default=True, type=str2bool)
      parser.add_argument('--wandb_project', default='ECCV2022_L2F', type=str)
      parser.add_argument('--wandb_table', default='Sample', type=str)
      parser.add_argument('--is_distributed', default=1, type=int, help='')
      parser.add_argument('--tasks', default=[1, 1, 1, 1, 1, 1, 1, 1, 0], type=str2list, help='task')
      parser.add_argument('--target_tasks', default=[0, 1, 0, 0, 0, 0, 1, 1, 0], type=str2list, help='target task')
      parser.add_argument('--lr_vnet', default=0.001, type=float, help='Vnet Learning Rate')
      parser.add_argument('--decay_vnet', default=0.0001, type=float, help='Vnet weight decay')
      parser.add_argument('--lr_meta', default=0.01, type=float, help='Meta Model Learning Rate(meta_lr)')
      parser.add_argument('--lr_meta_type', default=0, type=float, help='0:lr, 1:current lr')
      parser.add_argument('--activation', default='softmax', type=str)
      parser.add_argument('--activation_scaler', default='1.0', type=float)
      parser.add_argument('--t', default=1, help="temperature", type=float)
      parser.add_argument('--grad_clip', default=0, help="vnet gradient clip", type=float)
      parser.add_argument('--grad_clip_vnet', default=0, help="vnet gradient clip", type=float)
      parser.add_argument('--grad_clip_vnet_type', default=1, help="0:default, 1:ceil, 2:bottom", type=float)
      parser.add_argument('--ip_list', default=None, help="ipList", type=str2list)
      parser.add_argument('--loss_metric_align', default=2, help="", type=int)
      parser.add_argument('--loss_metric_joint', default=0, help="", type=int)
      parser.add_argument('--loss_type', default=1, help="0:all, 1:align + w*aux, 2:align_m_align + w*aux", type=int)
      parser.add_argument('--pri_type', default=0, help="0:pri in aux, 1:pri out aux", type=int)
      parser.add_argument('--vnet_split', default=0, help="0: shared, 1: split", type=int)
      parser.add_argument('--vnet_eval', default=1, help="0: train, 1: eval", type=int)
      parser.add_argument('--vnet_kind', default=1, help="0: MLP, 1:Transformer", type=int)
      parser.add_argument('--vnet_S_type', default=0, help="[l, t, l*t, (l,t)]", type=int)
      parser.add_argument('--vnet_T_type', default=3, help="[s, t, s*t, s+t, (s,t)]", type=int)
      parser.add_argument('--task_embed', default=100, help="", type=int)
      parser.add_argument('--warmup', default='warmup_linear', help="warmup_cosine, warmup_constant, warmup_linear, warmup_linear_down", type=str)
      parser.add_argument('--scheduler_epoch', default=4, help="", type=int)
      parser.add_argument('--milestones', default=None, help="", type=str2list)
      parser.add_argument('--milestones_proportion', default=0.1, help="", type=float)
      parser.add_argument("--warmup_proportion", default=0.1, type=float,
                          help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% of training.")
      parser.add_argument('--version', default="V1.0.0", type=str)
      parser.add_argument('--transformer_dim', default=[2048,128,128], type=str2list)
      parser.add_argument('--aux_lambda', default="1.0", type=float)
  
  
      parser.add_argument('--port', type=int, default=0, help='port')
      parser.add_argument('--weightnorm', type=str2bool, default=False, help='apply weight norm on layers')
      parser.add_argument('--aux_net', type=str, choices=['linear', 'nonlinear', 'transformer'], default='transformer')
      parser.add_argument('--nonlinear-init-upper', default=.2, type=float, help='upper value for nonlinear net')
      parser.add_argument('--skip', type=str2bool, default=False, help='skip connection on linear network')
      parser.add_argument('--main_task', type=int, default=1, help='')
      parser.add_argument('--auxgrad-every', type=int, default=1, help="number of opt. steps between aux params update")
      parser.add_argument('--hidden-dim', type=lambda s: [item.strip() for item in s.split(',')], default='10,10',
                          help="List of hidden dims for nonlinear")
  
      parser.add_argument('--eval_task', default=[1, 0], type=str2list, help="[retrieval, caption] : Evaluation Task")
      parser.add_argument('--eval_start', default=[0, 0], type=str2list, help='start evaluation')
      parser.add_argument('--eval_term', default=[1,1], type=str2list, help='start evaluation')
  
      parser.add_argument("--datatype", default="msrvtt7k", choices=['youcook', 'msrvtt1KA', 'msrvtt1KB', 'msrvtt7k', 'msrvttMIX'], type=str, help="Point the dataset `youcook` to finetune.")
      parser.add_argument('--BMM', default=[128, 48, 48], type=str2list, help="[BATCH, MAX_WORDS, MAX_FRAMES]")
      parser.add_argument('--debug', default=False, type=str2bool, help="")
      args = parser.parse_args() # auxgrad-every
      if args.port == 0:
          args.port = random.randint(0, 20000)
      if args.datatype == "youcook":
          args.n_display = 20
      else:
          args.n_display = 100
  
      if args.eval_term[0] == 0:
          args.eval_term[0] = 1
      if args.eval_term[1] == 0:
          args.eval_term[1] = 1
  
  
  
      args.batch_size, args.max_words, args.max_frames = args.BMM[0], args.BMM[1], args.BMM[2]
  
      # Check paramenters
      if args.gradient_accumulation_steps < 1:
          raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
              args.gradient_accumulation_steps))
      if not args.do_train and not args.do_eval:
          raise ValueError("At least one of `do_train` or `do_eval` must be True.")
  
      args.batch_size = int(args.batch_size / args.gradient_accumulation_steps)
  
      return args
  
  def set_seed_logger(args):
      global logger
  
      if args.seed:
          random.seed(args.seed)
          os.environ['PYTHONHASHSEED'] = str(args.seed)
          np.random.seed(args.seed)
          torch.manual_seed(args.seed)
          torch.cuda.manual_seed(args.seed)
          torch.cuda.manual_seed_all(args.seed)  # if you are using multi-GPU.
          torch.backends.cudnn.benchmark = False
          torch.backends.cudnn.deterministic = True
  
  
      if args.is_distributed:
          world_size = torch.distributed.get_world_size()
          torch.cuda.set_device(args.local_rank)
          args.world_size = world_size
  
      if not os.path.exists(args.output_dir):
          os.makedirs(args.output_dir, exist_ok=True)
  
      logger = get_logger(os.path.join(args.output_dir, "log.txt"))
  
      if args.local_rank == 0:
          logger.info("Effective parameters:")
          for key in sorted(args.__dict__):
              logger.info("  <<< {}: {}".format(key, args.__dict__[key]))
  
  
      return args
  
  def init_device(args, local_rank):
      global logger
  
      device = torch.device("cuda" if torch.cuda.is_available() else "cpu", local_rank)
  
      n_gpu = torch.cuda.device_count()
      logger.info("device: {} n_gpu: {}".format(device, n_gpu))
      args.n_gpu = n_gpu
  
      if args.batch_size % args.n_gpu != 0 or args.batch_size_val % args.n_gpu != 0:
          raise ValueError("Invalid batch_size/batch_size_val and n_gpu parameter: {}%{} and {}%{}, should be == 0".format(
              args.batch_size, args.n_gpu, args.batch_size_val, args.n_gpu))
  
      return device, n_gpu
  
  
  
  def init_model(args, device, n_gpu, local_rank):
  
      if args.init_model:
          model_state_dict = torch.load(args.init_model, map_location='cpu')
      else:
          model_state_dict = None
  
      # Prepare model
      cache_dir = args.cache_dir if args.cache_dir else os.path.join(str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed')
      model = UniVL.from_pretrained(args.bert_model, args.visual_model, args.cross_model, args.decoder_model,
                                     cache_dir=cache_dir, state_dict=model_state_dict, task_config=args)
  
      if args.is_distributed:
          model.to(device)
      else:
          model.cuda()
      return model
  
  def prep_optimizer(args, model, num_train_optimization_steps, num_train_warmup_steps, device, n_gpu, local_rank, coef_lr=1.):
  
      if hasattr(model, 'module'):
          model = model.module
  
      param_optimizer = list(model.named_parameters())
      no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
  
      no_decay_param_tp = [(n, p) for n, p in param_optimizer if not any(nd in n for nd in no_decay)]
      decay_param_tp = [(n, p) for n, p in param_optimizer if any(nd in n for nd in no_decay)]
  
      no_decay_bert_param_tp = [(n, p) for n, p in no_decay_param_tp if "bert." in n]
      no_decay_nobert_param_tp = [(n, p) for n, p in no_decay_param_tp if "bert." not in n]
  
      decay_bert_param_tp = [(n, p) for n, p in decay_param_tp if "bert." in n]
      decay_nobert_param_tp = [(n, p) for n, p in decay_param_tp if "bert." not in n]
  
      optimizer_grouped_parameters = [
          {'params': [p for n, p in no_decay_bert_param_tp], 'weight_decay': 0.01, 'lr': args.lr * coef_lr},
          {'params': [p for n, p in no_decay_nobert_param_tp], 'weight_decay': 0.01},
          {'params': [p for n, p in decay_bert_param_tp], 'weight_decay': 0.0, 'lr': args.lr * coef_lr},
          {'params': [p for n, p in decay_nobert_param_tp], 'weight_decay': 0.0}
      ]
  
      scheduler = None
      optimizer = BertAdam(optimizer_grouped_parameters, lr=args.lr, warmup=args.warmup, warmup_proportion=args.warmup_proportion,
                           milestones=args.milestones, milestones_proportion=args.milestones_proportion,
                           t_total=num_train_optimization_steps, t_warmup=num_train_warmup_steps, weight_decay=0.01,
                           max_grad_norm=1.0)
  
      if args.is_distributed:
          model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],
                                                            output_device=local_rank, find_unused_parameters=True)
  
      return optimizer, scheduler, model
  
  def dataloader_youcook_train(args, tokenizer):
      youcook_dataset = Youcook_Train(
          csv=args.youcook_train_csv,
          data_path=args.youcook_data_path,
          features_path=args.youcook_features_path,
          max_words=args.max_words,
          feature_framerate=args.feature_framerate,
          tokenizer=tokenizer,
          max_frames=args.max_frames,
      )
  
      if args.is_distributed:
          train_sampler = torch.utils.data.distributed.DistributedSampler(youcook_dataset)
      else:
          train_sampler = None
  
      dataloader = DataLoader(
          youcook_dataset,
          batch_size=args.batch_size // args.world_size,
          num_workers=args.num_thread_reader,
          pin_memory=False,
          shuffle=(train_sampler is None),
          sampler=train_sampler,
          drop_last=True,
      )
  
      return dataloader, len(youcook_dataset), train_sampler
  
  def dataloader_youcook_test(args, tokenizer):
      youcook_retrieval = Youcook_Retrieval_Test(
          csv=args.youcook_val_csv,
          data_path=args.youcook_data_path,
          features_path=args.youcook_features_path,
          max_words=args.max_words,
          feature_framerate=args.feature_framerate,
          tokenizer=tokenizer,
          max_frames=args.max_frames,
      )
  
  
      test_retrieval_sampler = SequentialSampler(youcook_retrieval)
      dataloader_youcook_retrieval = DataLoader(
          youcook_retrieval,
          sampler=test_retrieval_sampler,
          batch_size=args.batch_size_val,
          num_workers=args.num_thread_reader,
          pin_memory=False,
      )
  
      youcook_caption = Youcook_Caption_Test(
          csv=args.youcook_val_csv,
          data_path=args.youcook_data_path,
          features_path=args.youcook_features_path,
          max_words=args.max_words,
          feature_framerate=args.feature_framerate,
          tokenizer=tokenizer,
          max_frames=args.max_frames,
      )
  
      test_caption_sampler = SequentialSampler(youcook_caption)
      dataloader_youcook_caption = DataLoader(
          youcook_caption,
          sampler=test_caption_sampler,
          batch_size=args.batch_size_val,
          num_workers=args.num_thread_reader,
          pin_memory=False,
      )
  
      if args.local_rank == 0:
          logger.info('YoucookII retrieval validation pairs: {}'.format(len(youcook_retrieval)))
          logger.info('YoucookII caption validation pairs: {}'.format(len(youcook_caption)))
  
      return dataloader_youcook_retrieval, len(youcook_retrieval), dataloader_youcook_caption, len(youcook_caption)
  
  def dataloader_msrvtt_train(args, tokenizer):
      if args.datatype =="msrvtt1KA":
          msrvtt_dataset = MSRVTT_1KA_Train(
              csv_path=args.VTT_train_csv,
              json_path=args.VTT_data_path,
              features_path=args.VTT_features_path,
              max_words=args.max_words,
              feature_framerate=args.feature_framerate,
              tokenizer=tokenizer,
              max_frames=args.max_frames,
              unfold_sentences=args.expand_msrvtt_sentences,
          )
      elif args.datatype =="msrvtt1KB":
          msrvtt_dataset = MSRVTT_1KA_Train(
              csv_path=args.VTT_train_1KB_csv,
              json_path=args.VTT_data_path,
              features_path=args.VTT_features_path,
              max_words=args.max_words,
              feature_framerate=args.feature_framerate,
              tokenizer=tokenizer,
              max_frames=args.max_frames,
              unfold_sentences=args.expand_msrvtt_sentences,
          )
      elif args.datatype == "msrvtt7k":
          msrvtt_dataset = MSRVTT_7k_Train(
              csv_path=args.VTT_train_csv,
              json_path=args.VTT_data_path,
              features_path=args.VTT_features_path,
              max_words=args.max_words,
              feature_framerate=args.feature_framerate,
              tokenizer=tokenizer,
              max_frames=args.max_frames,
              split_type="train",
          )
      elif args.datatype == "msrvttMIX": # From 7K
          msrvtt_dataset = MSRVTT_MIX_Train(
              csv_path=args.VTT_train_csv,
              json_path=args.VTT_data_path,
              features_path=args.VTT_features_path,
              max_words=args.max_words,
              feature_framerate=args.feature_framerate,
              tokenizer=tokenizer,
              max_frames=args.max_frames,
              split_type="train",
          )
      else:
          raise argparse.ArgumentTypeError('Please Check Datatype')
  
      if args.is_distributed:
          train_sampler = torch.utils.data.distributed.DistributedSampler(msrvtt_dataset)
      else:
          train_sampler = None
  
      dataloader = DataLoader(
          msrvtt_dataset,
          batch_size=args.batch_size // args.world_size,
          num_workers=args.num_thread_reader,
          pin_memory=False,
          shuffle=(train_sampler is None),
          sampler=train_sampler,
          drop_last=True,
      )
  
      return dataloader, len(msrvtt_dataset), train_sampler
  
  
  
  
  def dataloader_msrvtt_test(args, tokenizer, split_type="test", ):
      if (args.datatype =="msrvtt1KA" or args.datatype == "msrvtt1KB") or args.datatype == "msrvttMIX":
          msrvtt_testset = MSRVTT_1KA_Test(
              csv_path=args.VTT_val_csv,
              features_path=args.VTT_features_path,
              max_words=args.max_words,
              feature_framerate=args.feature_framerate,
              tokenizer=tokenizer,
              max_frames=args.max_frames,
          )
      elif args.datatype == "msrvtt7k":
          msrvtt_testset = MSRVTT_7k_Test(
              csv_path=args.VTT_val_csv,
              json_path=args.VTT_data_path,
              features_path=args.VTT_features_path,
              max_words=args.max_words,
              feature_framerate=args.feature_framerate,
              tokenizer=tokenizer,
              max_frames=args.max_frames,
              split_type=split_type,
          )
      else:
          raise argparse.ArgumentTypeError('Please Check Datatype')
  
      test_msrvtt_sampler = SequentialSampler(msrvtt_testset)
      dataloader_msrvtt_test = DataLoader(
          msrvtt_testset,
          sampler=test_msrvtt_sampler,
          batch_size=args.batch_size_val,
          num_workers=args.num_thread_reader,
          pin_memory=False,
          drop_last=False,
      )
  
      return dataloader_msrvtt_test, len(msrvtt_testset), dataloader_msrvtt_test, len(msrvtt_testset)
  
  
  
  
  def save_model(epoch, args, model, type_name=""):
      # Only save the model it-self
      model_to_save = model.module if hasattr(model, 'module') else model
      output_model_file = os.path.join(
          args.output_dir, "pytorch_model.bin.{}{}".format("" if type_name=="" else type_name+".", epoch))
      torch.save(model_to_save.state_dict(), output_model_file)
      logger.info("Model saved to %s", output_model_file)
      return output_model_file
  
  def load_model(epoch, args, n_gpu, device, model_file=None):
      if model_file is None or len(model_file) == 0:
          model_file = os.path.join(args.output_dir, "pytorch_model.bin.{}".format(epoch))
      if os.path.exists(model_file):
          model_state_dict = torch.load(model_file, map_location='cpu')
          if args.local_rank == 0:
              logger.info("Model loaded from %s", model_file)
          # Prepare model
          cache_dir = args.cache_dir if args.cache_dir else os.path.join(str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed')
          model = UniVL.from_pretrained(args.bert_model, args.visual_model, args.cross_model, args.decoder_model,
                                         cache_dir=cache_dir, state_dict=model_state_dict, task_config=args)
  
          model.to(device)
      else:
          model = None
      return model
  def train_epoch(epoch, args, net, auxiliary_combine_net, train_dataloader, device, n_gpu, optimizer, meta_optimizer, scheduler, global_step, local_rank=0):
      global logger
      torch.cuda.empty_cache()
      net.train()
  
      log_step = args.n_display
  
  
      start_time_epoch = time.time()
      total_losses = AverageMeter()
  
      for step, batch in enumerate(tqdm(train_dataloader)):
          if n_gpu == 1 and args.is_distributed:
              batch = tuple(t.to(device=device, non_blocking=True) for t in batch)
          else:
              batch = tuple(t.cuda() for t in batch)
  
          input_ids, input_mask, segment_ids, video, video_mask, \
          pairs_masked_text, pairs_token_labels, masked_video, video_labels_index, \
          pairs_input_caption_ids, pairs_decoder_mask, pairs_output_caption_ids = batch
  
          losses = net(input_ids, segment_ids, input_mask, video, video_mask,
                       pairs_masked_text=pairs_masked_text, pairs_token_labels=pairs_token_labels,
                       masked_video=masked_video, video_labels_index=video_labels_index,
                       input_caption_ids=pairs_input_caption_ids, decoder_mask=pairs_decoder_mask,
                       output_caption_ids=pairs_output_caption_ids, tasks=args.tasks)
  
          loss, _ = auxiliary_combine_net(torch.stack(losses).unsqueeze(1), args.vnet_T_type)
  
          loss.backward()
          total_losses.update(float(loss.detach().cpu()))
          if args.is_distributed:
              torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
  
          if scheduler is not None:
              scheduler.step()  # Update learning rate schedule
  
          optimizer.step(epoch=epoch)
          optimizer.zero_grad()
  
          if (global_step+1) % args.auxgrad_every == 0:
              meta_val_loss = .0
              losses_pri = net(input_ids, segment_ids, input_mask, video, video_mask,
                                  pairs_masked_text=pairs_masked_text, pairs_token_labels=pairs_token_labels,
                                  masked_video=masked_video, video_labels_index=video_labels_index,
                                  input_caption_ids=pairs_input_caption_ids, decoder_mask=pairs_decoder_mask,
                                  output_caption_ids=pairs_output_caption_ids, tasks=args.target_tasks)
              meta_val_loss += sum(losses_pri)
  
              inner_loop_end_train_loss = 0.
              losses_train = net(input_ids, segment_ids, input_mask, video, video_mask,
                                  pairs_masked_text=pairs_masked_text, pairs_token_labels=pairs_token_labels,
                                  masked_video=masked_video, video_labels_index=video_labels_index,
                                  input_caption_ids=pairs_input_caption_ids, decoder_mask=pairs_decoder_mask,
                                  output_caption_ids=pairs_output_caption_ids, tasks=args.tasks)
              loss, _ = auxiliary_combine_net(torch.stack(losses_train).unsqueeze(1), args.vnet_T_type)
              inner_loop_end_train_loss += loss
  
              phi = list(auxiliary_combine_net.parameters())
              W = [p for n, p in net.named_parameters()]
  
              meta_optimizer.step(
                  val_loss=meta_val_loss,
                  train_loss=inner_loop_end_train_loss,
                  aux_params=phi,
                  parameters=W,
              )
  
          global_step += 1
          if global_step % log_step == 0 and local_rank == 0:
              logger.info("Epoch: %d/%s, Step: %d/%d, Lr: %s, Loss: %f", epoch + 1,
                          args.epochs, step + 1,
                          len(train_dataloader), "-".join([str('%.8f'%itm) for itm in sorted(list(set(optimizer.get_lr())))]),
                          total_losses.avg)
  
  
      return total_losses, global_step, time.time() - start_time_epoch
  
  def _run_on_single_gpu(net, batch_list_t, batch_list_v, batch_sequence_output_list, batch_visual_output_list):
      sim_matrix = []
      for idx1, b1 in enumerate(tqdm(batch_list_t)):
          input_ids, input_mask, segment_ids, _, _, _, _, _, _ = b1
          sequence_output = batch_sequence_output_list[idx1]
          each_row = []
          for idx2, b2 in enumerate(batch_list_v):
              _, _, _, video, video_mask, _, _, _, _ = b2
              visual_output = batch_visual_output_list[idx2]
              b1b2_logits = net.get_similarity_logits_align(sequence_output, visual_output, input_mask, video_mask)
              b1b2_logits = b1b2_logits.cpu().detach().numpy()
              each_row.append(b1b2_logits)
          each_row = np.concatenate(tuple(each_row), axis=-1)
          sim_matrix.append(each_row)
      return sim_matrix
  
  def eval_retrieval_epoch(args, model, test_dataloader, device, n_gpu):
  
      if hasattr(model, 'module'):
          model = model.module.to(device)
      else:
          model = model.to(device)
  
      if args.datatype != "youcook":
          test_dataloader.dataset.downstream = "retrieval"
  
      model.eval()
      with torch.no_grad():
          batch_list = []
          batch_sequence_output_list, batch_visual_output_list = [], []
          for bid, batch in enumerate(tqdm(test_dataloader)):
              batch = tuple(t.to(device) for t in batch)
  
              input_ids, input_mask, segment_ids, video, video_mask, _, _, _, _ = batch
  
  
              sequence_output, visual_output = model.get_sequence_visual_output(input_ids, segment_ids, input_mask, video, video_mask)
  
              batch_sequence_output_list.append(sequence_output)
              batch_visual_output_list.append(visual_output)
              batch_list.append(batch)
  
              print("{}/{}\r".format(bid, len(test_dataloader)), end="")
  
          if n_gpu > 1:
              device_ids = list(range(n_gpu))
              batch_list_t_splits = []
              batch_list_v_splits = []
              batch_t_output_splits = []
              batch_v_output_splits = []
              bacth_len = len(batch_list)
              split_len = (bacth_len + n_gpu - 1) // n_gpu
              for dev_id in device_ids:
                  s_, e_ = dev_id * split_len, (dev_id + 1) * split_len
                  if dev_id == 0:
                      batch_list_t_splits.append(batch_list[s_:e_])
                      batch_list_v_splits.append(batch_list)
  
                      batch_t_output_splits.append(batch_sequence_output_list[s_:e_])
                      batch_v_output_splits.append(batch_visual_output_list)
                  else:
                      devc = torch.device('cuda:{}'.format(str(dev_id)))
                      devc_batch_list = [tuple(t.to(devc) for t in b) for b in batch_list[s_:e_]]
                      batch_list_t_splits.append(devc_batch_list)
                      devc_batch_list = [tuple(t.to(devc) for t in b) for b in batch_list]
                      batch_list_v_splits.append(devc_batch_list)
  
                      devc_batch_list = [b.to(devc) for b in batch_sequence_output_list[s_:e_]]
                      batch_t_output_splits.append(devc_batch_list)
                      devc_batch_list = [b.to(devc) for b in batch_visual_output_list]
                      batch_v_output_splits.append(devc_batch_list)
              parameters_tuple_list = [(batch_list_t_splits[dev_id], batch_list_v_splits[dev_id],
                                        batch_t_output_splits[dev_id], batch_v_output_splits[dev_id]) for dev_id in device_ids]
              parallel_outputs = parallel_apply(_run_on_single_gpu, model, parameters_tuple_list, device_ids)
              sim_matrix = []
              for idx in range(len(parallel_outputs)):
                  sim_matrix += parallel_outputs[idx]
              sim_matrix = np.concatenate(tuple(sim_matrix), axis=0)
  
          else:
              sim_matrix = _run_on_single_gpu(model, batch_list, batch_list, batch_sequence_output_list, batch_visual_output_list)
  
      metrics = compute_metrics(sim_matrix)
      logger.info('\t Length-T: {}, Length-V:{}'.format(len(sim_matrix), len(sim_matrix[0])))
      logger.info('\t>>>  R@1: {:.4f} - R@5: {:.4f} - R@10: {:.4f} - Median R: {}'.
                  format(metrics['R1'], metrics['R5'], metrics['R10'], metrics['MR']))
  
      R1, R5, R10, MR = metrics['R1'], metrics['R5'], metrics['R10'], metrics['MR']
  
      return R1, R5, R10, MR
  
  
  
  
  
  DATALOADER_DICT = {}
  DATALOADER_DICT["youcook"] = {"train":dataloader_youcook_train, "val":dataloader_youcook_test}
  DATALOADER_DICT["msrvtt1KA"] = {"train":dataloader_msrvtt_train, "val":dataloader_msrvtt_test}
  DATALOADER_DICT["msrvtt7k"] = {"train":dataloader_msrvtt_train, "val":dataloader_msrvtt_test}
  DATALOADER_DICT["msrvttMIX"] = {"train":dataloader_msrvtt_train, "val":dataloader_msrvtt_test}
  
  def main_worker(gpu, ngpus_per_node, args):
      global logger
      args.local_rank = gpu  # 0
      args.world_rank = args.local_rank + args.world_rank * ngpus_per_node
      if args.is_distributed:
          print(args.ip, args.local_rank, args.world_rank, args.world_size)
          torch.distributed.init_process_group(backend="nccl", init_method='tcp://{}:{}'.format(args.ip, args.port), world_size=args.world_size, rank=args.world_rank)  # 8, 4
      args = set_seed_logger(args)
      if args.activation == "softmax":
          args.vnet_split = 0
  
  
      device, n_gpu = init_device(args, args.local_rank)
  
      tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
  
      assert  args.task_type == "retrieval"
      model = init_model(args, device, n_gpu, args.local_rank)
  
      if args.eval_task[1]:
          nlgEvalObj = NLGEval(no_overlap=False, no_skipthoughts=True, no_glove=True, metrics_to_omit=None)
      else:
          nlgEvalObj = None
  
      assert args.datatype in DATALOADER_DICT
      test_retrieval_dataloader, test_retrieval_length, test_caption_dataloader, test_caption_length = DATALOADER_DICT[args.datatype]["val"](args, tokenizer)
  
  
      if args.world_rank == 0 and args.use_wandb:
          wandb.init(project=args.wandb_project, name=args.wandb_table)
          args.dir = wandb.run.dir
          wandb.config.update(args)
  
      if args.local_rank == 0:
          logger.info("***** Running Retrieval test *****")
          logger.info("  Num examples = %d", test_retrieval_length)
          logger.info("  Batch size = %d", args.batch_size_val)
          logger.info("  Num steps = %d", len(test_retrieval_dataloader))
          logger.info("***** Running Captioning test *****")
          logger.info("  Num examples = %d", test_caption_length)
          logger.info("  Batch size = %d", args.batch_size_val)
          logger.info("  Num steps = %d", len(test_caption_dataloader))
  
      train_dataloader, train_length, train_sampler = DATALOADER_DICT[args.datatype]["train"](args, tokenizer)
      num_train_optimization_steps = (int(len(train_dataloader) + args.gradient_accumulation_steps - 1)
                                      / args.gradient_accumulation_steps) * args.epochs
      num_train_warmup_steps = (int(len(train_dataloader) + args.gradient_accumulation_steps - 1)
                                / args.gradient_accumulation_steps) * args.scheduler_epoch
  
      coef_lr = args.coef_lr
      if args.init_model:
          coef_lr = 1.0
      optimizer, scheduler, model = prep_optimizer(args, model, num_train_optimization_steps, num_train_warmup_steps, device, n_gpu, args.local_rank, coef_lr=coef_lr)
  
  
  
      # ================
      # hyperparam model
      # ================
      auxnet_mapping = dict(
          linear=MonoLinearHyperNet,
          nonlinear=MonoNonlinearHyperNet,
          transformer=MTN_auxi
      )
  
      auxnet_config = dict(t_dim=args.taskNum, f_dim = args.transformer_dim[0], i_dim = 1, h1_dim = args.transformer_dim[1], h2_dim = args.transformer_dim[2], o_dim=1)
  
      auxiliary_combine_net = auxnet_mapping[args.aux_net](**auxnet_config).to(device)
      auxiliary_param = list(auxiliary_combine_net.parameters())
  
      meta_opt = torch.optim.Adam(auxiliary_param, lr=args.lr_vnet, weight_decay=args.decay_vnet)
  
      meta_optimizer = MetaOptimizer(
          meta_optimizer=meta_opt, hpo_lr=args.lr, truncate_iter=3, max_grad_norm=50
      )
  
  
  
      if args.local_rank == 0:
          logger.info("***** Running training *****")
          logger.info("  Num examples = %d", train_length)
          logger.info("  Batch size = %d", args.batch_size)
          logger.info("  Num steps = %d", num_train_optimization_steps * args.gradient_accumulation_steps)
  
      best_score = [[0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0]]
      R1, R5, R10, MR = 0, 0, 0, 0
      Bleu_1, Bleu_2, Bleu_3, Bleu_4, METEOR, ROUGE_L, CIDEr = 0, 0, 0, 0, 0, 0, 0
      global_step = 0
  
      for epoch in range(args.epochs):
          if args.is_distributed:
              train_sampler.set_epoch(epoch)
  
          if args.is_distributed:
              train_sampler.set_epoch(epoch)
          if args.debug:
              tr_loss, global_step, ex_time = AverageMeter(), AverageMeter(), 0
          else:
              tr_loss, global_step, ex_time = train_epoch(epoch, args, model, auxiliary_combine_net, train_dataloader, device, n_gpu, optimizer, meta_optimizer,
                                                                                                          scheduler, global_step, local_rank=args.local_rank)
          if args.world_rank == 0:
              print("Epoch {}/{} Finished, Train Loss: {:0.4f}".format(epoch + 1, args.epochs, tr_loss.avg))
  
              if epoch >= args.eval_start[0] and (epoch % args.eval_term[0] == 0 and args.eval_task[0]):
                  R1, R5, R10, MR = eval_retrieval_epoch(args, model, test_retrieval_dataloader, device, n_gpu)
                  if best_score[0][0] <= R1:
                      best_score[0][0] = R1
                      best_score[0][1] = R5
                      best_score[0][2] = R10
                      best_score[0][3] = MR
  
                  print("Retrieval : Best R1 Accuracy : {:0.4f}, Current R1 Accuracy: {:0.4f}".format(best_score[0][0], R1))
              if epoch >= args.eval_start[1] and (epoch % args.eval_term[1] == 0 and args.eval_task[1]):
                  Bleu_1, Bleu_2, Bleu_3, Bleu_4, METEOR, ROUGE_L, CIDEr = eval_caption_epoch(args, model, test_caption_dataloader, tokenizer, device, n_gpu, logger, nlgEvalObj=nlgEvalObj)
  
                  if best_score[1][3] <= Bleu_4:
                      best_score[1][0] = Bleu_1
                      best_score[1][1] = Bleu_2
                      best_score[1][2] = Bleu_3
                      best_score[1][3] = Bleu_4
                      best_score[1][4] = METEOR
                      best_score[1][5] = ROUGE_L
                      best_score[1][6] = CIDEr
  
                  print("Caption : Best BLEU_4 Accuracy : {:0.4f}, Current BLEU_4 Accuracy: {:0.4f}".format(best_score[1][3], Bleu_4))
  
              current_lr = sorted(list(set(optimizer.get_lr())))[0] if len(sorted(list(set(optimizer.get_lr())))) != 0 else -1
              log_dict = {"train_loss": tr_loss.avg
                  , "R1": R1, "R5": R5, "R10": R10, "MR": MR
                  , "Best_R1": best_score[0][0], "Best_R5": best_score[0][1], "Best_R10": best_score[0][2], "Best_MR": best_score[0][3]
                  , "Bleu_1": Bleu_1, "Bleu_2": Bleu_2, "Bleu_3": Bleu_3, "Bleu_4": Bleu_4, "METEOR": METEOR, "ROUGE_L": ROUGE_L, "CIDEr": CIDEr
                  , "Best_Bleu_1": best_score[1][0], "Best_Bleu_2": best_score[1][1], "Best_Bleu_3": best_score[1][2], "Best_Bleu_4": best_score[1][3], "Best_METEOR": best_score[1][4], "Best_ROUGE_L": best_score[1][5], "Best_CIDEr": best_score[1][6]
                  , "current_lr": current_lr, "lr": args.lr, "ex_time_epoch": ex_time}
  
              if args.use_wandb:
                  wandb.log(log_dict)
                  torch.save(auxiliary_combine_net.state_dict(), os.path.join(wandb.run.dir, "MTN_{}.pt".format(epoch)))
                  torch.save({'args': args, }, os.path.join(wandb.run.dir, "args.pt"))
  
  
  
  def main():
      args = get_args()
      if args.is_distributed != 1:
          os.environ["CUDA_VISIBLE_DEVICES"] = "2"
      taskName = ["joint", "alignment", "mlm", "mfm", "m_joint", "m_align", "decoder", "m_decoder", "t_align", "t_align"]
      args.taskName = np.array(taskName)[np.array(args.tasks) == 1]
      args.taskNum = (np.array(args.tasks) == 1).sum()
      print("-"*10, args.taskName, "-"*10)
  
  
      s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
      s.connect(("8.8.8.8", 80))
      this_ip = s.getsockname()[0]
      s.close()
  
      if args.ip_list: # multi-node
          ip_list = args.ip_list
          args.ip = ip_list[0]
      else: # single-node
          args.ip = this_ip
          ip_list = [args.ip]
  
      args.world_size = len(ip_list)  # 8
      for i, ip in enumerate(ip_list):
          if ip == this_ip:
              args.world_rank = i
  
      ngpus_per_node = torch.cuda.device_count()  # 8
      args.world_size = ngpus_per_node * args.world_size  # 8 * 8
      if args.is_distributed:
          print(args.world_size)
          mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args)) # args.world_size or ngpus_per_node
      else:
          main_worker(0, ngpus_per_node, args)
  
  if __name__ == "__main__":
      main()