text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import pdb
  import torch
  import torch.nn as nn
  import torch.optim as optim
  
  import numpy as np
  
  import time
  import os
  from six.moves import cPickle
  
  import opts
  import models
  from dataloader_123k import *
  import eval_utils
  import misc.utils as utils
  from misc.rewards import init_scorer, get_self_critical_reward
  
  try:
      import tensorboardX as tb
  except ImportError:
      print("tensorboardX is not installed")
      tb = None
  
  def add_summary_value(writer, key, value, iteration):
      if writer:
          writer.add_scalar(key, value, iteration)
  
  def train(opt):
      # Deal with feature things before anything
      opt.use_att = utils.if_use_att(opt.caption_model)
      if opt.use_box: opt.att_feat_size = opt.att_feat_size + 5
  
      loader = DataLoader(opt)
      opt.vocab_size = loader.vocab_size
      opt.seq_length = loader.seq_length
      
      tb_summary_writer = tb and tb.SummaryWriter(opt.checkpoint_path)
  
      infos = {}
      histories = {}
      if opt.start_from is not None:
          # open old infos and check if models are compatible
          with open(os.path.join(opt.start_from, 'infos_'+opt.id+'.pkl')) as f:
              infos = cPickle.load(f)
              saved_model_opt = infos['opt']
              need_be_same=["rnn_type", "rnn_size", "num_layers"] #removed caption_model temporarily
              for checkme in need_be_same:
                  assert vars(saved_model_opt)[checkme] == vars(opt)[checkme], "Command line argument and saved model disagree on '%s' " % checkme
  
          if os.path.isfile(os.path.join(opt.start_from, 'histories_'+opt.id+'.pkl')):
              with open(os.path.join(opt.start_from, 'histories_'+opt.id+'.pkl')) as f:
                  histories = cPickle.load(f)
  
      iteration = infos.get('iter', 0)
      epoch = infos.get('epoch', 0)
  
      val_result_history = histories.get('val_result_history', {})
      loss_history = histories.get('loss_history', {})
      lr_history = histories.get('lr_history', {})
      ss_prob_history = histories.get('ss_prob_history', {})
  
      loader.iterators = infos.get('iterators', loader.iterators)
      loader.split_ix = infos.get('split_ix', loader.split_ix)
      if opt.load_best_score == 1:
          best_val_score = infos.get('best_val_score', None)
  
      model = models.setup(opt).cuda()
      
      dp_model = torch.nn.DataParallel(model)
  
      update_lr_flag = True
      # Assure in training mode
      dp_model.train()
      
      crit = utils.LanguageModelCriterion()
      rl_crit = utils.RewardCriterion()
      if opt.reduce_on_plateau:
          optimizer = utils.build_optimizer(model.parameters(), opt)
          optimizer = utils.ReduceLROnPlateau(optimizer, factor=0.5, patience=3)
      else:    
          optimizer = utils.build_optimizer(model.parameters(), opt)
      # Load the optimizer
      if vars(opt).get('start_from', None) is not None and os.path.isfile(os.path.join(opt.start_from,"optimizer.pth")):
          optimizer.load_state_dict(torch.load(os.path.join(opt.start_from, 'optimizer.pth')))
  # model.load_state_dict(torch.load(opt.model, map_location=lambda storage, loc: storage))
      while True:
          if update_lr_flag:
              if not opt.reduce_on_plateau:
                  # Assign the learning rate
                  if epoch > opt.learning_rate_decay_start and opt.learning_rate_decay_start >= 0:
                      frac = (epoch - opt.learning_rate_decay_start) // opt.learning_rate_decay_every
                      decay_factor = opt.learning_rate_decay_rate  ** frac
                      opt.current_lr = opt.learning_rate * decay_factor
                  else:
                      opt.current_lr = opt.learning_rate
                  utils.set_lr(optimizer, opt.current_lr)
                  # Assign the scheduled sampling prob
              if epoch > opt.scheduled_sampling_start and opt.scheduled_sampling_start >= 0:
                  frac = (epoch - opt.scheduled_sampling_start) // opt.scheduled_sampling_increase_every
                  opt.ss_prob = min(opt.scheduled_sampling_increase_prob  * frac, opt.scheduled_sampling_max_prob)
                  model.ss_prob = opt.ss_prob
  
              # If start self critical training
              if opt.self_critical_after != -1 and epoch >= opt.self_critical_after:
                  sc_flag = True
                  init_scorer(opt.cached_tokens)
              else:
                  sc_flag = False
  
              update_lr_flag = False
                  
          start = time.time()
          # Load data from train split (0)
          data = loader.get_batch('train')
  #         print('Read data:', time.time() - start)
          '''checking data object'''
  
          
          torch.cuda.synchronize()
          start = time.time()
  
          tmp = [data['fc_feats'], data['att_feats'], data['labels'], data['masks'], data['att_masks']]
          tmp = [_ if _ is None else torch.from_numpy(_).cuda() for _ in tmp]
          fc_feats, att_feats, labels, masks, att_masks = tmp
          
          optimizer.zero_grad()
          if not sc_flag:
              loss = crit(dp_model(fc_feats, att_feats, labels, att_masks), labels[:,1:], masks[:,1:])
              
          else:
              gen_result, sample_logprobs = dp_model(fc_feats, att_feats, att_masks, opt={'sample_max':0}, mode='sample')
              reward = get_self_critical_reward(dp_model, fc_feats, att_feats, att_masks, data, gen_result, opt)
              loss = rl_crit(sample_logprobs, gen_result.data, torch.from_numpy(reward).float().cuda())
  
          loss.backward()
          utils.clip_gradient(optimizer, opt.grad_clip)
          optimizer.step()
          train_loss = loss.item()
          torch.cuda.synchronize()
          end = time.time()
          if not sc_flag:
  #             pass
              print("iter {} (epoch {}), train_loss = {:.3f}, time/batch = {:.3f}" \
                  .format(iteration, epoch, train_loss, end - start))
          else:
  #             pass
              print("iter {} (epoch {}), avg_reward = {:.3f}, time/batch = {:.3f}" \
                  .format(iteration, epoch, np.mean(reward[:,0]), end - start))
  
          # Update the iteration and epoch
          iteration += 1
          if data['bounds']['wrapped']:
              epoch += 1
              update_lr_flag = True
              print("Iter: {}, Loss: {}, Epoc: {}, LR: {}".format(iteration, train_loss, epoch, opt.current_lr))
  
          # Write the training loss summary
          if (iteration % opt.losses_log_every == 0):
              add_summary_value(tb_summary_writer, 'train_loss', train_loss, iteration)
              if opt.reduce_on_plateau:
                  opt.current_lr = optimizer.current_lr
              add_summary_value(tb_summary_writer, 'learning_rate', opt.current_lr, iteration)
              add_summary_value(tb_summary_writer, 'scheduled_sampling_prob', model.ss_prob, iteration)
              if sc_flag:
                  add_summary_value(tb_summary_writer, 'avg_reward', np.mean(reward[:,0]), iteration)
  
              loss_history[iteration] = train_loss if not sc_flag else np.mean(reward[:,0])
              lr_history[iteration] = opt.current_lr
              ss_prob_history[iteration] = model.ss_prob
  
          # make evaluation on validation set, and save model
          if (iteration % opt.save_checkpoint_every == 0):
              # eval model
              eval_kwargs = {'split': 'val',
                              'dataset': opt.input_json}
              eval_kwargs.update(vars(opt))
              val_loss, predictions, lang_stats = eval_utils.eval_split(dp_model, crit, loader, eval_kwargs)
              if opt.reduce_on_plateau:
                  if 'CIDEr' in lang_stats:
                      optimizer.scheduler_step(-lang_stats['CIDEr'])
                  else:
                      optimizer.scheduler_step(val_loss)
                      
              # Write validation result into summary
              add_summary_value(tb_summary_writer, 'validation loss', val_loss, iteration)
              if lang_stats is not None:
                  for k,v in lang_stats.items():
                      add_summary_value(tb_summary_writer, k, v, iteration)
              val_result_history[iteration] = {'loss': val_loss, 'lang_stats': lang_stats, 'predictions': predictions}
  
              # Save model if is improving on validation result
              if opt.language_eval == 1:
                  current_score = lang_stats['CIDEr']
              else:
                  current_score = - val_loss
  
              best_flag = False
              if True: # if true
                  if best_val_score is None or current_score > best_val_score:
                      best_val_score = current_score
                      best_flag = True
                  checkpoint_path = os.path.join(opt.checkpoint_path, 'model.pth')
                  torch.save(model.state_dict(), checkpoint_path)
                  print("model saved to {}".format(checkpoint_path))
                  optimizer_path = os.path.join(opt.checkpoint_path, 'optimizer.pth')
                  torch.save(optimizer.state_dict(), optimizer_path)
  
                  # Dump miscalleous informations
                  infos['iter'] = iteration
                  infos['epoch'] = epoch
                  infos['iterators'] = loader.iterators
                  infos['split_ix'] = loader.split_ix
                  infos['best_val_score'] = best_val_score
                  infos['opt'] = opt
                  infos['vocab'] = loader.get_vocab()
  
                  histories['val_result_history'] = val_result_history
                  histories['loss_history'] = loss_history
                  histories['lr_history'] = lr_history
                  histories['ss_prob_history'] = ss_prob_history
                  with open(os.path.join(opt.checkpoint_path, 'infos_'+opt.id+'.pkl'), 'wb') as f:
                      cPickle.dump(infos, f)
                  with open(os.path.join(opt.checkpoint_path, 'histories_'+opt.id+'.pkl'), 'wb') as f:
                      cPickle.dump(histories, f)
  
                  if best_flag:
                      checkpoint_path = os.path.join(opt.checkpoint_path, 'model-best.pth')
                      torch.save(model.state_dict(), checkpoint_path)
                      print("model saved to {}".format(checkpoint_path))
                      with open(os.path.join(opt.checkpoint_path, 'infos_'+opt.id+'-best.pkl'), 'wb') as f:
                          cPickle.dump(infos, f)
  
          # Stop if reaching max epochs
          if epoch >= opt.max_epochs and opt.max_epochs != -1:
              break
  
  opt = opts.parse_opt()
  train(opt)
  