text: import argparse
  import itertools
  import math 
  import os 
  import logging 
  
  from tensorboardX import SummaryWriter
  import torch
  from torch import nn, optim
  from torch.optim import lr_scheduler
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  import yaml
  from bisect import bisect
  from pathlib import Path 
  
  from visdialch.data.dataset import VisDialDataset
  from visdialch.data.bert_dataset import BertVisDialDataset
  from visdialch.encoders import Encoder
  from visdialch.decoders import Decoder
  from visdialch.metrics import SparseGTMetrics, NDCG
  from visdialch.model import EncoderDecoderModel
  from visdialch.utils.checkpointing import CheckpointManager, load_checkpoint
  from visdialch.loss import NpairLoss
  
  from transformers import AdamW, WarmupCosineWithHardRestartsSchedule
  from transformers import WarmupCosineSchedule, WarmupLinearSchedule 
  
  import pdb
  
  
  parser = argparse.ArgumentParser()
  parser.add_argument(
      "--config-yml",
      default="configs/lf_disc_faster_rcnn_x101.yml",
      help="Path to a config file listing reader, model and solver parameters.",
  )
  parser.add_argument(
      "--train-json",
      default="data/visdial_1.0_train.json",
      help="Path to json file containing VisDial v1.0 training data.",
  )
  parser.add_argument(
      "--val-json",
      default="data/visdial_1.0_val.json",
      help="Path to json file containing VisDial v1.0 validation data.",
  )
  parser.add_argument(
      "--val-dense-json",
      default="data/visdial_1.0_val_dense_annotations.json",
      help="Path to json file containing VisDial v1.0 validation dense ground "
      "truth annotations.",
  )
  
  
  parser.add_argument_group(
      "Arguments independent of experiment reproducibility"
  )
  parser.add_argument(
      "--gpu-ids",
      nargs="+",
      type=int,
      default=0,
      help="List of ids of GPUs to use.",
  )
  parser.add_argument(
      "--cpu-workers",
      type=int,
      default=4,
      help="Number of CPU workers for dataloader.",
  )
  parser.add_argument(
      "--overfit",
      action="store_true",
      help="Overfit model on 5 examples, meant for debugging.",
  )
  parser.add_argument(
      "--validate",
      action="store_true",
      help="Whether to validate on val split after every epoch.",
  )
  parser.add_argument(
      "--in-memory",
      action="store_true",
      help="Load the whole dataset and pre-extracted image features in memory. "
      "Use only in presence of large RAM, atleast few tens of GBs.",
  )
  
  
  parser.add_argument_group("Checkpointing related arguments")
  parser.add_argument(
      "--save-dirpath",
      default="checkpoints/",
      help="Path of directory to create checkpoint directory and save "
      "checkpoints.",
  )
  parser.add_argument(
      "--load-pthpath",
      default="",
      help="To continue training, path to .pth file of saved checkpoint.",
  )
  
  # For reproducibility.
  # Refer https://pytorch.org/docs/stable/notes/randomness.html
  torch.manual_seed(0)
  torch.cuda.manual_seed_all(0)
  torch.backends.cudnn.benchmark = False
  torch.backends.cudnn.deterministic = True
  
  
  # =============================================================================
  #   INPUT ARGUMENTS AND CONFIG
  # =============================================================================
  
  args = parser.parse_args()
  
  def prepare_logger(dirpath):
      Path(dirpath).mkdir(parents=True, exist_ok=True)
      # Create a custom logger
      logger = logging.getLogger(__name__)
      logger.setLevel(logging.INFO)
  
      # Create handlers
      c_handler = logging.StreamHandler()
      f_handler = logging.FileHandler(os.path.join(dirpath, 'log.txt'), mode='w')
  
      # Create formatters and add it to handlers
      c_format = logging.Formatter('%(message)s')
      f_format = logging.Formatter('%(message)s')
      c_handler.setFormatter(c_format)
      f_handler.setFormatter(f_format)
  
      # Add handlers to the logger
      logger.addHandler(c_handler)
      logger.addHandler(f_handler)
  
      return logger
  
  logger = prepare_logger(args.save_dirpath)
  
  # keys: {"dataset", "model", "solver"}
  config = yaml.load(open(args.config_yml))
  
  if isinstance(args.gpu_ids, int):
      args.gpu_ids = [args.gpu_ids]
  device = (
      torch.device("cuda", args.gpu_ids[0])
      if args.gpu_ids[0] >= 0
      else torch.device("cpu")
  )
  
  # Print config and args.
  logger.info(yaml.dump(config, default_flow_style=False))
  for arg in vars(args):
      logger.info("{:<20}: {}".format(arg, getattr(args, arg)))
  
  
  # =============================================================================
  #   SETUP DATASET, DATALOADER, MODEL, CRITERION, OPTIMIZER, SCHEDULER
  # =============================================================================
  
  word_embedding_type = config['dataset']['word_embedding_type']
  if word_embedding_type not in ['init', 'glove', 'bert']:
      raise NotImplementedError()
  logger.info('Word embedding type: {}'.format(word_embedding_type))
  
  if word_embedding_type == 'bert':
      train_dataset = BertVisDialDataset(
          config["dataset"],
          args.train_json,
          overfit=args.overfit,
          in_memory=args.in_memory,
          return_options=config["model"]["decoder"] == "disc",
          add_boundary_toks=config["model"]["decoder"] == "gen",
          proj_to_senq_id=config["model"]["decoder"] == "gen"
      )
  
      val_dataset = BertVisDialDataset(
          config["dataset"],
          args.val_json,
          args.val_dense_json,
          overfit=args.overfit,
          in_memory=args.in_memory,
          return_options=True,
          add_boundary_toks=config["model"]["decoder"] == "gen",
          proj_to_senq_id=config["model"]["decoder"] == "gen"
      )
  else:
      train_dataset = VisDialDataset(
          config["dataset"],
          args.train_json,
          overfit=args.overfit,
          in_memory=args.in_memory,
          return_options=config["model"]["decoder"] == "disc",
          add_boundary_toks=config["model"]["decoder"] == "gen"
      )
  
      val_dataset = VisDialDataset(
          config["dataset"],
          args.val_json,
          args.val_dense_json,
          overfit=args.overfit,
          in_memory=args.in_memory,
          return_options=True,
          add_boundary_toks=config["model"]["decoder"] == "gen"
      )
  
  assert((config["solver"]['batch_size'] 
          % config["solver"]["accumulation_steps"]) == 0)
  actual_batch_size = (config["solver"]['batch_size'] 
                       // config["solver"]["accumulation_steps"])
  logger.info('Actual batch size: {}'.format(actual_batch_size))
  logger.info('Gradient accumulation steps: {}'.format(config['solver']['accumulation_steps']))
  logger.info('Effective batch size: {}'.format(config['solver']['batch_size']))
  
  train_dataloader = DataLoader(
      train_dataset,
      batch_size=actual_batch_size,
      num_workers=args.cpu_workers,
      shuffle=True,
  )
  
  val_dataloader = DataLoader(
      val_dataset,
      batch_size=actual_batch_size
      if config["model"]["decoder"] == "disc"
      else len(args.gpu_ids),
      num_workers=args.cpu_workers,
  )
  
  
  # Pass vocabulary to construct Embedding layer.
  encoder = Encoder(config["model"], train_dataset.vocabulary)
  if word_embedding_type == 'bert':
      decoder = Decoder(
          config["model"], train_dataset.vocabulary, 
          bert_model=encoder.word_embed.bert
      )
  else:
      decoder = Decoder(config["model"], train_dataset.vocabulary)
  logger.info("Encoder: {}".format(config["model"]["encoder"]))
  logger.info("Decoder: {}".format(config["model"]["decoder"]))
  
  # Share word embedding between encoder and decoder.
  if not word_embedding_type == 'bert':
      decoder.word_embed = encoder.word_embed
  
  # Wrap encoder and decoder in a model.
  model = EncoderDecoderModel(encoder, decoder).to(device)
  if -1 not in args.gpu_ids:
      model = nn.DataParallel(model, args.gpu_ids)
  
  # Loss function.
  if config["model"]["decoder"] == "disc":
      if config["model"]["loss"] == "ce":
         criterion = nn.CrossEntropyLoss()
      elif config["model"]["loss"] == "np":
         criterion = NpairLoss(scale=0.25)
      else:
         raise NotImplementedError
  elif config["model"]["decoder"] == "gen":
      criterion = nn.CrossEntropyLoss(
          ignore_index=(
              model.module.decoder.padding_idx
              if isinstance(model, nn.DataParallel)
              else model.decoder.padding_idx
          )
      )
  else:
      raise NotImplementedError
  
  if config["solver"]["training_splits"] == "trainval":
      data_iterations = len(train_dataloader) + len(val_dataloader)
  else:
      data_iterations = len(train_dataloader)
  iterations = int(math.ceil(data_iterations / config['solver']['accumulation_steps']))
  
  
  def lr_lambda_fun(current_iteration: int) -> float:
      """Returns a learning rate multiplier.
  
      Till `warmup_epochs`, learning rate linearly increases to `initial_lr`,
      and then gets multiplied by `lr_gamma` every time a milestone is crossed.
      """
      current_epoch = float(current_iteration) / iterations
      if current_epoch <= config["solver"]["warmup_epochs"]:
          alpha = current_epoch / float(config["solver"]["warmup_epochs"])
          return config["solver"]["warmup_factor"] * (1.0 - alpha) + alpha
      else:
          idx = bisect(config["solver"]["lr_milestones"], current_epoch)
          return pow(config["solver"]["lr_gamma"], idx)
  
  
  reason_mode = config['model']['reason_mode']
  logger.info('Reason mode: {}'.format(reason_mode))
  
  bert_params = [item for item in list(model.named_parameters()) if 'bert' in item[0]]
  not_bert_params = [item for item in list(model.named_parameters()) if 'bert' not in item[0]]
  
  if 'no_decay' in config['solver'].keys():
      no_decay = config['solver']['no_decay']
  else:
      no_decay = []
  
  grouped_parameters = [
      {
          'params': [p for n, p in not_bert_params if not any(nd in n for nd in no_decay)], 
          'weight_decay': config['solver']['weight_decay']
      }, 
      {
          'params': [p for n, p in bert_params if not any(nd in n for nd in no_decay)],
          'weight_decay': 0.01
      },
      {
          'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 
          'weight_decay': 0.0
      }
  ]
  
  if config['solver']['optimizer'] == 'Adamax':
      optimizer = optim.Adamax(
          grouped_parameters, 
          lr=config["solver"]["initial_lr"],
          weight_decay=config['solver']['weight_decay']
      )
  elif config['solver']['optimizer'] == 'AdamW':
      optimizer = AdamW(
          grouped_parameters, 
          lr=config["solver"]["initial_lr"], 
          weight_decay=config['solver']['weight_decay']
      )
  else:
      raise NotImplementedError()
  logger.info('Optimizer: {}'.format(config['solver']['optimizer']))
      
  if config['solver']['lr_schedule'] == 'warmup_reduce_on_milestones':
      scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda_fun)
  elif config['solver']['lr_schedule'] == 'warmup_linear':
      warmup_steps = iterations * config["solver"]["warmup_epochs"]
      t_total = iterations * config["solver"]["num_epochs"]
      scheduler = WarmupLinearSchedule(
          optimizer, warmup_steps=warmup_steps, t_total=t_total
      )
  elif config['solver']['lr_schedule'] == 'warmup_cosine':
      warmup_steps = iterations * config["solver"]["warmup_epochs"]
      t_total = iterations * config["solver"]["num_epochs"]
      scheduler = WarmupCosineSchedule(
          optimizer, warmup_steps=warmup_steps, t_total=t_total
      )
  elif config['solver']['lr_schedule'] == 'warmup_cosine_with_hard_restarts':
      warmup_steps = iterations * config["solver"]["warmup_epochs"]
      t_total = iterations * config["solver"]["num_epochs"]
      cycles = config['solver']['cycles']
      scheduler = WarmupCosineWithHardRestartsSchedule(
          optimizer, warmup_steps=warmup_steps, t_total=t_total, cycles=cycles
      )
  else:
      raise NotImplementedError()
  logger.info('Learning rate schedule: {}'.format(config['solver']['lr_schedule']))
  
  
  # =============================================================================
  #   SETUP BEFORE TRAINING LOOP
  # =============================================================================
  
  summary_writer = SummaryWriter(log_dir=args.save_dirpath)
  checkpoint_manager = CheckpointManager(
      model, optimizer, args.save_dirpath, config=config
  )
  sparse_metrics = SparseGTMetrics()
  ndcg = NDCG()
  
  # If loading from checkpoint, adjust start epoch and load parameters.
  if args.load_pthpath == "":
      start_epoch = 0
  else:
      # "path/to/checkpoint_xx.pth" -> xx
      start_epoch = int(args.load_pthpath.split("_")[-1][:-4])
  
      model_state_dict, optimizer_state_dict = load_checkpoint(args.load_pthpath)
      if isinstance(model, nn.DataParallel):
          model.module.load_state_dict(model_state_dict)
      else:
          model.load_state_dict(model_state_dict)
      optimizer.load_state_dict(optimizer_state_dict)
      logger.info("Loaded model from {}".format(args.load_pthpath))
  
  # =============================================================================
  #   TRAINING LOOP
  # =============================================================================
  
  # Forever increasing counter to keep track of iterations (for tensorboard log).
  global_iteration_step = start_epoch * iterations
  model.train()   # Explicitly set the model to train mode 
  
  for epoch in range(start_epoch, config["solver"]["num_epochs"]):
  
      # -------------------------------------------------------------------------
      #   ON EPOCH START  (combine dataloaders if training on train + val)
      # -------------------------------------------------------------------------
      if config["solver"]["training_splits"] == "trainval":
          combined_dataloader = itertools.chain(train_dataloader, val_dataloader)
      else:
          combined_dataloader = itertools.chain(train_dataloader)
  
      logger.info(f"\nTraining for epoch {epoch}:")
      #pdb.set_trace()
      optimizer.zero_grad()
      for i, batch in enumerate(tqdm(combined_dataloader)):
          for key in batch:
              batch[key] = batch[key].to(device)
  
          output = model(batch)
          target = (
              batch["ans_ind"]
              if config["model"]["decoder"] == "disc"
              else batch["ans_out"]
          )
          batch_loss = criterion(
              output.view(-1, output.size(-1)), target.view(-1)
          )
          batch_loss.backward()
  
          if ((i+1) % config['solver']['accumulation_steps'] == 0
              or (i+1) == data_iterations):
              if config['solver']['max_grad_norm'] > 0:
                  torch.nn.utils.clip_grad_norm_(model.parameters(), 
                                                 config["solver"]["max_grad_norm"])
              optimizer.step()
              optimizer.zero_grad()
  
              summary_writer.add_scalar(
                  "train/loss", batch_loss, global_iteration_step
              )
              summary_writer.add_scalar(
                  "train/lr", optimizer.param_groups[0]["lr"], global_iteration_step
              )
  
              scheduler.step(global_iteration_step)
              global_iteration_step += 1
              torch.cuda.empty_cache()
  
      # -------------------------------------------------------------------------
      #   ON EPOCH END  (checkpointing and validation)
      # -------------------------------------------------------------------------
      checkpoint_manager.step()
  
      # Validate and report automatic metrics.
      if args.validate:
  
          # Switch dropout, batchnorm etc to the correct mode.
          model.eval()
  
          logger.info(f"\nValidation after epoch {epoch}:")
          for i, batch in enumerate(tqdm(val_dataloader)):
              for key in batch:
                  batch[key] = batch[key].to(device)
              with torch.no_grad():
                  output = model(batch)
              sparse_metrics.observe(output, batch["ans_ind"])
              if "gt_relevance" in batch:
                  output = output[
                      torch.arange(output.size(0)), batch["round_id"] - 1, :
                  ]
                  ndcg.observe(output, batch["gt_relevance"])
  
          all_metrics = {}
          all_metrics.update(sparse_metrics.retrieve(reset=True))
          all_metrics.update(ndcg.retrieve(reset=True))
          for metric_name, metric_value in all_metrics.items():
              logger.info(f"{metric_name}: {metric_value}")
          summary_writer.add_scalars(
              "metrics", all_metrics, global_iteration_step
          )
  
          model.train()
          torch.cuda.empty_cache()
          summary_writer.flush()
          
  summary_writer.close()