text: from __future__ import division
  
  import sys, tempfile
  import onmt
  import onmt.markdown
  import onmt.modules
  import argparse
  import torch
  import torch.nn as nn
  from torch import cuda
  from torch.autograd import Variable
  import math
  import time, datetime
  import random
  import numpy as np
  from onmt.multiprocessing.multiprocessing_wrapper import MultiprocessingRunner
  from onmt.train_utils.trainer import BaseTrainer
  
  
  
  class AETrainer(BaseTrainer):
  
      def __init__(self, autoencoder, model, loss_function, trainData, validData, dicts, opt):
          super().__init__(model, loss_function, trainData, validData, dicts, opt)
          self.optim = onmt.Optim(opt)
          self.autoencoder = autoencoder
          if(opt.auto_encoder_type is None):
              self.auto_encoder_type = "Baseline"
          else:
              self.auto_encoder_type = opt.auto_encoder_type
  
          if self.cuda:
              torch.cuda.set_device(self.opt.gpus[0])
              torch.manual_seed(self.opt.seed)
              self.loss_function = self.loss_function.cuda()
              self.model = self.model.cuda()
              self.autoencoder = self.autoencoder.cuda()
  
          self.optim.set_parameters(self.autoencoder.parameters())
  
      def save(self, epoch, valid_ppl, batchOrder=None, iteration=-1):
  
          opt = self.opt
          autoencoder = self.autoencoder
  
          autoencoder_state_dict = self.autoencoder.state_dict()
          optim_state_dict = self.optim.state_dict()
  
          #  drop a checkpoint
          checkpoint = {
              'autoencoder': autoencoder_state_dict,
              'opt': opt,
              'epoch': epoch,
              'iteration': iteration,
              'batchOrder': batchOrder,
              'optim': optim_state_dict
          }
  
          file_name = '%s_ppl_%.2f_e%.2f.pt' % (opt.save_model, valid_ppl, epoch)
          print('Writing to %s' % file_name)
          torch.save(checkpoint, file_name)
  
          # check te save directory here
  
      def eval(self, data):
          total_loss = 0
          total_words = 0
  
          batch_order = data.create_order(random=False)
          self.model.eval()
          self.autoencoder.eval()
          """ New semantics of PyTorch: save space by not creating gradients """
          with torch.no_grad():
              for i in range(len(data)):
  
                  batch = data.next()[0]
                  if (self.cuda):
                      batch.cuda()
  
                  """ outputs can be either 
                          hidden states from decoder or
                          prob distribution from decoder generator
                  """
  
                  targets,outputs = self.autoencoder(batch)
  
                  loss_data = self.loss_function(outputs, targets)
  
                  # ~
                  total_loss += loss_data
                  total_words += outputs.size(0)
  
          self.autoencoder.train()
          return total_loss / total_words
  
      def train_epoch(self, epoch, resume=False, batchOrder=None, iteration=0):
  
          opt = self.opt
          train_data = self.train_data
  
          # Clear the gradients of the model
          # self.runner.zero_grad()
          self.autoencoder.zero_grad()
          self.model.eval()
  
          if opt.extra_shuffle and epoch > opt.curriculum:
              train_data.shuffle()
  
          # Shuffle mini batch order.
  
          if resume:
              train_data.batchOrder = batchOrder
              train_data.set_index(iteration)
              print("Resuming from iteration: %d" % iteration)
          else:
              batchOrder = train_data.create_order()
              iteration = 0
  
          total_loss, total_words = 0, 0
          report_loss, report_mu,report_sig,report_el,report_mse, report_kl, report_tgt_words = 0, 0,0,0,0,0,0
          report_src_words = 0
          start = time.time()
          nSamples = len(train_data)
  
          counter = 0
          num_accumulated_words = 0
          num_accumulated_sents = 0
  
          for i in range(iteration, nSamples):
  
              curriculum = (epoch < opt.curriculum)
  
              batch = train_data.next(curriculum=curriculum)[0]
              if (self.cuda):
                  batch.cuda()
  
              oom = False
              try:
  
                  batch_size = batch.size
  
                  # print("Input size:",batch[0].size())
                  targets,outputs = self.autoencoder(batch)
  
                  loss_data= self.loss_function(outputs, targets.data)
                  if(self.auto_encoder_type == "Variational"):
                      m = self.autoencoder.variational_layer.mean
                      std = self.autoencoder.variational_layer.std
                      m = m.mul(m)
                      one = torch.ones(m.size())
                      if(m.is_cuda):
                          one = one.cuda()
                      var_loss = ((m+std-std.log()-one)*0.5).sum()
                      report_mse += loss_data.item()
                      report_kl += var_loss.item()
                      report_mu += m.sum().item()
                      report_sig += std.sum().item()
                      report_el += m.numel()
                      loss_data = loss_data + var_loss
                  loss_data.backward()
  
  
              except RuntimeError as e:
                  if 'out of memory' in str(e):
                      print('| WARNING: ran out of memory on GPU , skipping batch')
                      oom = True
                      torch.cuda.empty_cache()
                  else:
                      raise e
  
              if not oom:
  
                  counter = counter + 1
                  num_accumulated_words += targets.size(0)
                  num_accumulated_sents += batch_size
  
  
                  if num_accumulated_words >= opt.batch_size_update * 0.95:
                      grad_denom = 1
                      if self.opt.normalize_gradient:
                          grad_denom = num_accumulated_words
                      # Update the parameters.
                      self.optim.step(grad_denom=grad_denom)
                      self.autoencoder.zero_grad()
                      self.model.zero_grad()
                      counter = 0
                      num_accumulated_words = 0
                      num_accumulated_sents = 0
                      num_updates = self.optim._step
  
                      if opt.save_every > 0 and num_updates % opt.save_every == -1 % opt.save_every:
                          valid_loss = self.eval(self.valid_data)
                          print('Validation perplexity: %g' % valid_loss)
  
                          ep = float(epoch) - 1. + ((float(i) + 1.) / nSamples)
  
                          self.save(ep, valid_loss, batchOrder=batchOrder, iteration=i)
  
                  num_words = targets.size(0)
                  report_loss += loss_data
                  report_tgt_words += num_words
                  total_loss += loss_data
                  total_words += num_words
  
                  optim = self.optim
  
                  if i == 0 or (i % opt.log_interval == -1 % opt.log_interval):
                      print(("Epoch %2d, %5d/%5d; ; loss: %6.2f (%6.2f, %6.2f) ; var: mu %6.2f sig: %6.2f; lr: %.7f ; num updates: %7d " +
                             "%5.0f src tok/s; %s elapsed") %
                            (epoch, i + 1, len(train_data),
                             report_loss / report_tgt_words,report_mse/report_tgt_words,report_kl/report_tgt_words,
                             report_mu / max(1,report_el), report_sig / max(1,report_el),
                             optim.getLearningRate(),
                             optim._step,
                             report_tgt_words / (time.time() - start),
                             str(datetime.timedelta(seconds=int(time.time() - self.start_time)))))
                      report_loss, report_tgt_words ,report_mse,report_kl= 0, 0, 0,0
                      report_mu,report_sig,report_el = 0,0,0
                      report_src_words = 0
                      start = time.time()
  
          return total_loss / total_words
  
      def run(self, save_file=None):
  
          opt = self.opt
          model = self.model
          autoencoder = self.autoencoder
          optim = self.optim
  
          # Try to load the save_file
          batchOrder = None
          iteration = 0
          print('Initializing model parameters')
          autoencoder.init_model_parameters()
          resume = False
  
          valid_loss = self.eval(self.valid_data)
          print('Validation loss: %g' % valid_loss)
  
          self.start_time = time.time()
  
          for epoch in range(opt.start_epoch, opt.start_epoch + opt.epochs):
              print('')
  
              #  (1) train for one epoch on the training set
              train_loss = self.train_epoch(epoch, resume=resume,
                                            batchOrder=batchOrder,
                                            iteration=iteration)
              print('Train loss: %g' % train_loss)
  
              #  (2) evaluate on the validation set
              valid_loss = self.eval(self.valid_data)
              print('Validation perplexity: %g' % valid_loss)
  
              self.save(epoch, valid_loss)
              batchOrder = None
              iteration = None
              resume = False
  
  
  
  
  
  