text: from __future__ import print_function
  from six.moves import range
  
  import torch
  import torch.nn as nn
  import torch.optim as optim
  from torch.autograd import Variable
  import torch.backends.cudnn as cudnn
  
  from PIL import Image
  
  from miscc.config import cfg
  from miscc.utils import mkdir_p
  from miscc.utils import build_super_images, build_super_images2
  from miscc.utils import weights_init, load_params, copy_G_params
  from model import G_DCGAN, G_NET
  from datasets import prepare_data
  from model import RNN_ENCODER, CNN_ENCODER
  
  # from miscc.losses import words_loss
  # from miscc.losses import discriminator_loss, generator_loss, KL_loss
  from rsr_losses import words_loss
  from rsr_losses import discriminator_loss, generator_loss, KL_loss, swd_loss
  import os
  import time
  import numpy as np
  import sys
  
  import gc
  
  # ################# Text to image task############################ #
  class condGANTrainer(object):
      def __init__(self, output_dir, data_loader, n_words, ixtoword):
          if cfg.TRAIN.FLAG:
              self.model_dir = os.path.join(output_dir, 'Model')
              self.image_dir = os.path.join(output_dir, 'Image')
              mkdir_p(self.model_dir)
              mkdir_p(self.image_dir)
  
          torch.cuda.set_device(cfg.GPU_ID)
          cudnn.benchmark = True
  
          self.batch_size = cfg.TRAIN.BATCH_SIZE
          self.max_epoch = cfg.TRAIN.MAX_EPOCH
          self.snapshot_interval = cfg.TRAIN.SNAPSHOT_INTERVAL
  
          self.n_words = n_words
          self.ixtoword = ixtoword
          self.data_loader = data_loader
          self.num_batches = len(self.data_loader)
  
      def _show_resource(self):
          for obj in gc.get_objects():
              try:
                  if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
                      print(type(obj), obj.size())
              except:
                  pass
  
      def build_models(self):
          # ###################encoders######################################## #
          if cfg.TRAIN.NET_E == '':
              print('Error: no pretrained text-image encoders')
              return
  
          image_encoder = CNN_ENCODER(cfg.TEXT.EMBEDDING_DIM)
          img_encoder_path = cfg.TRAIN.NET_E.replace('text_encoder', 'image_encoder')
          state_dict = \
              torch.load(img_encoder_path, map_location=lambda storage, loc: storage)
          image_encoder.load_state_dict(state_dict)
          for p in image_encoder.parameters():
              p.requires_grad = False
          print('Load image encoder from:', img_encoder_path)
          image_encoder.eval()
  
          text_encoder = \
              RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
          state_dict = \
              torch.load(cfg.TRAIN.NET_E,
                         map_location=lambda storage, loc: storage)
          text_encoder.load_state_dict(state_dict)
          for p in text_encoder.parameters():
              p.requires_grad = False
          print('Load text encoder from:', cfg.TRAIN.NET_E)
          text_encoder.eval()
  
          # #######################generator and discriminators############## #
          netsD = []
          if cfg.GAN.B_DCGAN:
              if cfg.TREE.BRANCH_NUM ==1:
                  from model import D_NET64 as D_NET
              elif cfg.TREE.BRANCH_NUM == 2:
                  from model import D_NET128 as D_NET
              else:  # cfg.TREE.BRANCH_NUM == 3:
                  from model import D_NET256 as D_NET
              # TODO: elif cfg.TREE.BRANCH_NUM > 3:
              netG = G_DCGAN()
              netsD = [D_NET(b_jcu=False)]
          else:
              from model import D_NET64, D_NET128, D_NET256
              netG = G_NET()
              if cfg.TREE.BRANCH_NUM > 0:
                  netsD.append(D_NET64())
              if cfg.TREE.BRANCH_NUM > 1:
                  netsD.append(D_NET128())
              if cfg.TREE.BRANCH_NUM > 2:
                  netsD.append(D_NET256())
              # TODO: if cfg.TREE.BRANCH_NUM > 3:
          netG.apply(weights_init)
          # print(netG)
          for i in range(len(netsD)):
              netsD[i].apply(weights_init)
              # print(netsD[i])
          print('# of netsD', len(netsD))
          #
          epoch = 0
          if cfg.TRAIN.NET_G != '':
              state_dict = \
                  torch.load(cfg.TRAIN.NET_G, map_location=lambda storage, loc: storage)
              netG.load_state_dict(state_dict)
              print('Load G from: ', cfg.TRAIN.NET_G)
              istart = cfg.TRAIN.NET_G.rfind('_') + 1
              iend = cfg.TRAIN.NET_G.rfind('.')
              epoch = cfg.TRAIN.NET_G[istart:iend]
              epoch = int(epoch) + 1
              if cfg.TRAIN.B_NET_D:
                  Gname = cfg.TRAIN.NET_G
                  for i in range(len(netsD)):
                      s_tmp = Gname[:Gname.rfind('/')]
                      Dname = '%s/netD%d.pth' % (s_tmp, i)
                      print('Load D from: ', Dname)
                      state_dict = \
                          torch.load(Dname, map_location=lambda storage, loc: storage)
                      netsD[i].load_state_dict(state_dict)
          # ########################################################### #
          if cfg.CUDA:
              text_encoder = text_encoder.cuda()
              image_encoder = image_encoder.cuda()
              netG.cuda()
              for i in range(len(netsD)):
                  netsD[i].cuda()
          return [text_encoder, image_encoder, netG, netsD, epoch]
  
      def define_optimizers(self, netG, netsD):
          optimizersD = []
          num_Ds = len(netsD)
          for i in range(num_Ds):
              opt = optim.Adam(netsD[i].parameters(),
                               lr=cfg.TRAIN.DISCRIMINATOR_LR,
                               betas=(0.5, 0.999))
              optimizersD.append(opt)
  
          optimizerG = optim.Adam(netG.parameters(),
                                  lr=cfg.TRAIN.GENERATOR_LR,
                                  betas=(0.5, 0.999))
  
          return optimizerG, optimizersD
  
      def prepare_labels(self):
          batch_size = self.batch_size
          real_labels = Variable(torch.FloatTensor(batch_size).fill_(1))
          fake_labels = Variable(torch.FloatTensor(batch_size).fill_(0))
          match_labels = Variable(torch.LongTensor(range(batch_size)))
          if cfg.CUDA:
              real_labels = real_labels.cuda()
              fake_labels = fake_labels.cuda()
              match_labels = match_labels.cuda()
  
          return real_labels, fake_labels, match_labels
  
      def save_model(self, netG, avg_param_G, netsD, epoch):
          backup_para = copy_G_params(netG)
          load_params(netG, avg_param_G)
          torch.save(netG.state_dict(),
              '%s/netG_epoch_%d.pth' % (self.model_dir, epoch))
          load_params(netG, backup_para)
          #
          for i in range(len(netsD)):
              netD = netsD[i]
              torch.save(netD.state_dict(),
                  '%s/netD%d.pth' % (self.model_dir, i))
          print('Save G/Ds models.')
  
      def set_requires_grad_value(self, models_list, brequires):
          for i in range(len(models_list)):
              for p in models_list[i].parameters():
                  p.requires_grad = brequires
  
      def save_img_results(self, netG, noise, sent_emb, words_embs, mask,
                           image_encoder, captions, cap_lens,
                           gen_iterations, name='current'):
          # Save images
          fake_imgs, attention_maps, _, _ = netG(noise, sent_emb, words_embs, mask)
          for i in range(len(attention_maps)):
              if len(fake_imgs) > 1:
                  img = fake_imgs[i + 1].detach().cpu()
                  lr_img = fake_imgs[i].detach().cpu()
              else:
                  img = fake_imgs[0].detach().cpu()
                  lr_img = None
              attn_maps = attention_maps[i]
              att_sze = attn_maps.size(2)
              img_set, _ = \
                  build_super_images(img, captions, self.ixtoword,
                                     attn_maps, att_sze, lr_imgs=lr_img)
              if img_set is not None:
                  im = Image.fromarray(img_set)
                  fullpath = '%s/G_%s_%d_%d.png'\
                      % (self.image_dir, name, gen_iterations, i)
                  im.save(fullpath)
  
          # for i in range(len(netsD)):
          i = -1
          img = fake_imgs[i].detach()
          region_features, _ = image_encoder(img)
          att_sze = region_features.size(2)
          _, _, att_maps = words_loss(region_features.detach(),
                                      words_embs.detach(),
                                      None, cap_lens,
                                      None, self.batch_size)
          img_set, _ = \
              build_super_images(fake_imgs[i].detach().cpu(),
                                 captions, self.ixtoword, att_maps, att_sze)
          if img_set is not None:
              im = Image.fromarray(img_set)
              fullpath = '%s/D_%s_%d.png'\
                  % (self.image_dir, name, gen_iterations)
              im.save(fullpath)
  
      def rsr_train(self):
          text_encoder, image_encoder, netG, netsD, start_epoch = self.build_models()
          avg_param_G = copy_G_params(netG)
          optimizerG, optimizersD = self.define_optimizers(netG, netsD)
          real_labels, fake_labels, match_labels = self.prepare_labels()
  
          batch_size = self.batch_size
          nz = cfg.GAN.Z_DIM
          # parameters of rsr
          num_small_batch = cfg.TRAIN.NUM_BATCH_SIZE
          d_img = 8192
          N_rotmat = 100
          large_batch_size = num_small_batch * batch_size
          # noise = Variable(torch.FloatTensor(batch_size, nz))
          noise = Variable(torch.FloatTensor(large_batch_size, nz))
          all_real_features = torch.zeros(cfg.TREE.BRANCH_NUM, large_batch_size, d_img)
          all_fake_features = torch.zeros(cfg.TREE.BRANCH_NUM, large_batch_size, d_img)
          fixed_noise = Variable(torch.FloatTensor(batch_size, nz).normal_(0, 1))
          if cfg.CUDA:
              noise, fixed_noise = noise.cuda(), fixed_noise.cuda()
              all_real_features, all_fake_features = all_real_features.cuda(), all_fake_features.cuda()
  
          gen_iterations = 0
          # gen_iterations = start_epoch * self.num_batches
          for epoch in range(start_epoch, self.max_epoch):
              start_t = time.time()
  
              data_iter = iter(self.data_loader)
              step = 0
              while step < self.num_batches:
                  # reset requires_grad to be trainable for all Ds
                  # self.set_requires_grad_value(netsD, True)
  
                  ######################################################
                  # (1) Prepare training data and Compute text embeddings
                  ######################################################
                  data = data_iter.next()
                  # TODO: check what "class_ids" represents
                  imgs, captions, cap_lens, class_ids, keys = prepare_data(data)
  
                  # hidden = text_encoder.init_hidden(batch_size)
                  hidden = text_encoder.init_hidden(large_batch_size)
                  # words_embs: batch_size x nef x seq_len
                  # sent_emb: batch_size x nef
                  words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
                  words_embs, sent_emb = words_embs.detach(), sent_emb.detach()
                  mask = (captions == 0)
                  num_words = words_embs.size(2)
                  if mask.size(1) > num_words:
                      mask = mask[:, :num_words]
  
                  #######################################################
                  # (2) Generate noise of fake images
                  ######################################################
                  # based on noise + text embeddings
                  noise.data.normal_(0, 1)
                  all_fake_imgs = []
                  with torch.no_grad():
                      for j in range(num_small_batch):
                          fake_imgs, _, _, _ = netG(noise[j * batch_size: (j + 1) * batch_size],
                                                    sent_emb[j * batch_size: (j + 1) * batch_size],
                                                    words_embs[j * batch_size: (j + 1) * batch_size],
                                                    mask[j * batch_size: (j + 1) * batch_size])
                          all_fake_imgs.append(fake_imgs)
                  #######################################################
                  # (3) Update D network
                  ######################################################
                  # DEBUG
                  # print("before dis")
                  # os.system("nvidia-smi")
                  # self._show_resource()
  
                  errD_total = 0
                  D_logs = ''
                  for j in range(num_small_batch):
                      # fake_imgs, _, _, _ = netG(noise[j * batch_size: (j + 1) * batch_size],
                      #                                 sent_emb[j * batch_size: (j + 1) * batch_size],
                      #                                 words_embs[j * batch_size: (j + 1) * batch_size],
                      #                                 mask[j * batch_size: (j + 1) * batch_size])
                      # sliced_imgs = [arr[j * batch_size: (j + 1) * batch_size] for arr in imgs]
                      errD_total_small_batch = 0
                      for i in range(len(netsD)):
                          netsD[i].zero_grad()
                          errD = discriminator_loss(netsD[i],
                                                    imgs[i][j * batch_size: (j + 1) * batch_size],
                                                    all_fake_imgs[j][i],
                                                    sent_emb[j * batch_size: (j + 1) * batch_size],
                                                    real_labels, fake_labels)
                          # backward and update parameters
                          errD.backward()
                          errD_total_small_batch += errD.detach()
                          D_logs += 'errD%d: %.2f ' % (i, errD.data[0])
                          optimizersD[i].step()
                      errD_total += errD_total_small_batch
  
                      # del fake_imgs
                      # del sliced_imgs
                      # torch.cuda.empty_cache()
  
                  #######################################################
                  # (4) Update G network: maximize log(D(G(z)))
                  ######################################################
                  # compute total loss for training G
                  step += 1
                  gen_iterations += 1
  
                  # DEBUG
                  # print("before RSR-run")
                  # os.system("nvidia-smi")
                  # self._show_resource()
  
                  ### RUN
                  with torch.no_grad():
                      for j in range(num_small_batch):
                          # fake_imgs, _, _, _ = netG(noise[j * batch_size: (j + 1) * batch_size],
                          #                                 sent_emb[j * batch_size: (j + 1) * batch_size],
                          #                                 words_embs[j * batch_size: (j + 1) * batch_size],
                          #                                 mask[j * batch_size: (j + 1) * batch_size])
                          # sliced_imgs = [arr[j * batch_size: (j + 1) * batch_size] for arr in imgs]
                          # DEBUG
                          # print("before RSR-run-beforeassign")
                          # os.system("nvidia-smi")
                          for i in range(len(netsD)):
                              # all_real_features[i][j * batch_size: (j + 1) * batch_size] = netsD[i](sliced_imgs[i]).view(batch_size, -1)
                              # all_fake_features[i][j * batch_size: (j + 1) * batch_size] = netsD[i](fake_imgs[i]).view(batch_size, -1)
                              all_real_features[i][j * batch_size: (j + 1) * batch_size].copy_(netsD[i](imgs[i][j * batch_size: (j + 1) * batch_size]).view(
                                  batch_size, -1))
                              all_fake_features[i][j * batch_size: (j + 1) * batch_size].copy_(netsD[i](all_fake_imgs[j][i]).view(
                                  batch_size, -1))
  
                  # del all_fake_imgs
                  # del sliced_imgs
                  # torch.cuda.empty_cache()
  
                  # DEBUG
                  # print("before RSR-sort")
                  # os.system("nvidia-smi")
                  # self._show_resource()
  
                  ### SORT
                  rotmat_img = torch.randn(d_img, N_rotmat)
                  if cfg.CUDA:
                      rotmat_img = rotmat_img.cuda()
  
                  all_out_img_sort_relative = []
                  with torch.no_grad():
                      for i in range(len(netsD)):
                          all_real_features_projected = all_real_features[i].mm(rotmat_img)
                          all_fake_features_projected = all_fake_features[i].mm(rotmat_img)
  
                          # [_, out_img_sort_ix] = torch.sort(all_fake_features_projected, dim=0)
                          # [_, out_img_sort_relative] = out_img_sort_ix.sort(0)
                          # all_out_img_sort_relative.append(out_img_sort_relative)
                          # [all_real_features_projected_sorted, _] = torch.sort(all_real_features_projected, dim=0)
  
                          [_, out_img_sort_ix] = torch.sort(all_fake_features_projected.cpu(), dim=0)
                          [_, out_img_sort_relative] = out_img_sort_ix.sort(0)
                          all_out_img_sort_relative.append(out_img_sort_relative.cuda())
  
                          [all_real_features_projected_sorted, _] = torch.sort(all_real_features_projected.cpu(), dim=0)
                          all_real_features_projected_sorted = all_real_features_projected_sorted.cuda()
                      # del all_real_features_projected
                      # del all_fake_features_projected
                      # del out_img_sort_ix
                      # torch.cuda.empty_cache()
  
                  # DEBUG
                  # print("before RSR-rerun")
                  # os.system("nvidia-smi")
                  # self._show_resource()
  
                  ### RERUN
                  # do not need to compute gradient for Ds
                  # self.set_requires_grad_value(netsD, False)
                  netG.zero_grad()
                  # errG_total, G_logs = \
                  #     generator_loss(netsD, image_encoder, fake_imgs, real_labels,
                  #                    words_embs, sent_emb, match_labels, cap_lens, class_ids)
                  G_logs_total = ''
                  for j in range(num_small_batch):
                      fake_imgs, _, mu, logvar = netG(noise[j * batch_size: (j + 1) * batch_size],
                                                      sent_emb[j * batch_size: (j + 1) * batch_size],
                                                      words_embs[j * batch_size: (j + 1) * batch_size],
                                                      mask[j * batch_size: (j + 1) * batch_size])
                      # sliced_imgs = [arr[j * batch_size: (j + 1) * batch_size] for arr in imgs]
                      # print(all_out_img_sort_relative[i][j * batch_size: (j + 1) * batch_size].shape)
                      # print(imgs[i].shape)
                      # sliced_imgs = [imgs[i].gather(0, all_out_img_sort_relative[i][j * batch_size: (j + 1) * batch_size]) for i in range(len(imgs))]
                      sliced_real_features = [all_real_features_projected_sorted.gather(0, all_out_img_sort_relative[i][j * batch_size: (j + 1) * batch_size]) for i in range(len(imgs))]
                      errG_total, G_logs, fake_features, fake_projected, real_sorted, fake_sorted = \
                          swd_loss(netsD, image_encoder, fake_imgs, real_labels,
                                   words_embs[j * batch_size: (j + 1) * batch_size],
                                   sent_emb[j * batch_size: (j + 1) * batch_size],
                                   match_labels, cap_lens,
                                   class_ids[j * batch_size: (j + 1) * batch_size],
                                   sliced_real_features, rotmat_img)
  
                      kl_loss = KL_loss(mu, logvar)
                      errG_total += kl_loss
                      G_logs += 'kl_loss: %.2f ' % kl_loss.data[0]
                      G_logs_total += G_logs
                      # backward and update parameters
                      errG_total.backward()
  
                  optimizerG.step()
  
                  del all_real_features_projected_sorted
                  del fake_imgs
                  del sliced_real_features
                  del rotmat_img
                  del fake_features, fake_projected, real_sorted, fake_sorted
                  del all_out_img_sort_relative
                  torch.cuda.empty_cache()
  
                  # DEBUG
                  # print("after RSR")
                  # self._show_resource()
                  # os.system("nvidia-smi")
                  # if step == 2:
                  #     exit()
  
                  for p, avg_p in zip(netG.parameters(), avg_param_G):
                      avg_p.mul_(0.999).add_(0.001, p.data)
  
                  if gen_iterations % 100 == 0:
                      print(D_logs + '\n' + G_logs)
                  # save images
                  # if gen_iterations % 1000 == 0:
                  #     backup_para = copy_G_params(netG)
                  #     load_params(netG, avg_param_G)
                  #     self.save_img_results(netG, fixed_noise, sent_emb,
                  #                           words_embs, mask, image_encoder,
                  #                           captions, cap_lens, epoch, name='average')
                  #     load_params(netG, backup_para)
                      #
                      # self.save_img_results(netG, fixed_noise, sent_emb,
                      #                       words_embs, mask, image_encoder,
                      #                       captions, cap_lens,
                      #                       epoch, name='current')
              end_t = time.time()
  
              print('''[%d/%d][%d]
                            Loss_D: %.2f Loss_G: %.2f Time: %.2fs'''
                    % (epoch, self.max_epoch, self.num_batches,
                       errD_total.data[0], errG_total.data[0],
                       end_t - start_t))
  
              del errG_total, errD_total
              torch.cuda.empty_cache()
  
              if epoch % cfg.TRAIN.SNAPSHOT_INTERVAL == 0:  # and epoch != 0:
                  self.save_model(netG, avg_param_G, netsD, epoch)
  
          self.save_model(netG, avg_param_G, netsD, self.max_epoch)
  
      def train(self):
          text_encoder, image_encoder, netG, netsD, start_epoch = self.build_models()
          avg_param_G = copy_G_params(netG)
          optimizerG, optimizersD = self.define_optimizers(netG, netsD)
          real_labels, fake_labels, match_labels = self.prepare_labels()
  
          batch_size = self.batch_size
          nz = cfg.GAN.Z_DIM
          noise = Variable(torch.FloatTensor(batch_size, nz))
          fixed_noise = Variable(torch.FloatTensor(batch_size, nz).normal_(0, 1))
          if cfg.CUDA:
              noise, fixed_noise = noise.cuda(), fixed_noise.cuda()
  
          gen_iterations = 0
          # gen_iterations = start_epoch * self.num_batches
          for epoch in range(start_epoch, self.max_epoch):
              start_t = time.time()
  
              data_iter = iter(self.data_loader)
              step = 0
              while step < self.num_batches:
                  # reset requires_grad to be trainable for all Ds
                  # self.set_requires_grad_value(netsD, True)
  
                  ######################################################
                  # (1) Prepare training data and Compute text embeddings
                  ######################################################
                  data = data_iter.next()
                  imgs, captions, cap_lens, class_ids, keys = prepare_data(data)
  
                  hidden = text_encoder.init_hidden(batch_size)
                  # words_embs: batch_size x nef x seq_len
                  # sent_emb: batch_size x nef
                  words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
                  words_embs, sent_emb = words_embs.detach(), sent_emb.detach()
                  mask = (captions == 0)
                  num_words = words_embs.size(2)
                  if mask.size(1) > num_words:
                      mask = mask[:, :num_words]
  
                  #######################################################
                  # (2) Generate fake images
                  ######################################################
                  noise.data.normal_(0, 1)
                  fake_imgs, _, mu, logvar = netG(noise, sent_emb, words_embs, mask)
  
                  #######################################################
                  # (3) Update D network
                  ######################################################
                  errD_total = 0
                  D_logs = ''
                  for i in range(len(netsD)):
                      netsD[i].zero_grad()
                      errD = discriminator_loss(netsD[i], imgs[i], fake_imgs[i],
                                                sent_emb, real_labels, fake_labels)
                      # backward and update parameters
                      errD.backward()
                      optimizersD[i].step()
                      errD_total += errD
                      D_logs += 'errD%d: %.2f ' % (i, errD.data[0])
  
                  #######################################################
                  # (4) Update G network: maximize log(D(G(z)))
                  ######################################################
                  # compute total loss for training G
                  step += 1
                  gen_iterations += 1
  
                  # do not need to compute gradient for Ds
                  # self.set_requires_grad_value(netsD, False)
                  netG.zero_grad()
                  errG_total, G_logs = \
                      generator_loss(netsD, image_encoder, fake_imgs, real_labels,
                                     words_embs, sent_emb, match_labels, cap_lens, class_ids)
                  kl_loss = KL_loss(mu, logvar)
                  errG_total += kl_loss
                  G_logs += 'kl_loss: %.2f ' % kl_loss.data[0]
                  # backward and update parameters
                  errG_total.backward()
                  optimizerG.step()
                  for p, avg_p in zip(netG.parameters(), avg_param_G):
                      avg_p.mul_(0.999).add_(0.001, p.data)
  
                  if gen_iterations % 100 == 0:
                      print(D_logs + '\n' + G_logs)
                  # save images
                  if gen_iterations % 1000 == 0:
                      backup_para = copy_G_params(netG)
                      load_params(netG, avg_param_G)
                      self.save_img_results(netG, fixed_noise, sent_emb,
                                            words_embs, mask, image_encoder,
                                            captions, cap_lens, epoch, name='average')
                      load_params(netG, backup_para)
                      #
                      # self.save_img_results(netG, fixed_noise, sent_emb,
                      #                       words_embs, mask, image_encoder,
                      #                       captions, cap_lens,
                      #                       epoch, name='current')
              end_t = time.time()
  
              print('''[%d/%d][%d]
                    Loss_D: %.2f Loss_G: %.2f Time: %.2fs'''
                    % (epoch, self.max_epoch, self.num_batches,
                       errD_total.data[0], errG_total.data[0],
                       end_t - start_t))
  
              if epoch % cfg.TRAIN.SNAPSHOT_INTERVAL == 0:  # and epoch != 0:
                  self.save_model(netG, avg_param_G, netsD, epoch)
  
          self.save_model(netG, avg_param_G, netsD, self.max_epoch)
  
      def save_singleimages(self, images, filenames, save_dir,
                            split_dir, sentenceID=0):
          for i in range(images.size(0)):
              s_tmp = '%s/single_samples/%s/%s' %\
                  (save_dir, split_dir, filenames[i])
              folder = s_tmp[:s_tmp.rfind('/')]
              if not os.path.isdir(folder):
                  print('Make a new folder: ', folder)
                  mkdir_p(folder)
  
              fullpath = '%s_%d.jpg' % (s_tmp, sentenceID)
              # range from [-1, 1] to [0, 1]
              # img = (images[i] + 1.0) / 2
              img = images[i].add(1).div(2).mul(255).clamp(0, 255).byte()
              # range from [0, 1] to [0, 255]
              ndarr = img.permute(1, 2, 0).data.cpu().numpy()
              im = Image.fromarray(ndarr)
              im.save(fullpath)
  
      def sampling(self, split_dir):
          if cfg.TRAIN.NET_G == '':
              print('Error: the path for morels is not found!')
          else:
              if split_dir == 'test':
                  split_dir = 'valid'
              # Build and load the generator
              if cfg.GAN.B_DCGAN:
                  netG = G_DCGAN()
              else:
                  netG = G_NET()
              netG.apply(weights_init)
              netG.cuda()
              netG.eval()
              #
              text_encoder = RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
              state_dict = \
                  torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
              text_encoder.load_state_dict(state_dict)
              print('Load text encoder from:', cfg.TRAIN.NET_E)
              text_encoder = text_encoder.cuda()
              text_encoder.eval()
  
              batch_size = self.batch_size
              nz = cfg.GAN.Z_DIM
              noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
              noise = noise.cuda()
  
              model_dir = cfg.TRAIN.NET_G
              state_dict = \
                  torch.load(model_dir, map_location=lambda storage, loc: storage)
              # state_dict = torch.load(cfg.TRAIN.NET_G)
              netG.load_state_dict(state_dict)
              print('Load G from: ', model_dir)
  
              # the path to save generated images
              s_tmp = model_dir[:model_dir.rfind('.pth')]
              save_dir = '%s/%s' % (s_tmp, split_dir)
              mkdir_p(save_dir)
  
              cnt = 0
  
              for _ in range(1):  # (cfg.TEXT.CAPTIONS_PER_IMAGE):
                  for step, data in enumerate(self.data_loader, 0):
                      cnt += batch_size
                      if step % 100 == 0:
                          print('step: ', step)
                      # if step > 50:
                      #     break
  
                      imgs, captions, cap_lens, class_ids, keys = prepare_data(data)
  
                      hidden = text_encoder.init_hidden(batch_size)
                      # words_embs: batch_size x nef x seq_len
                      # sent_emb: batch_size x nef
                      words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
                      words_embs, sent_emb = words_embs.detach(), sent_emb.detach()
                      mask = (captions == 0)
                      num_words = words_embs.size(2)
                      if mask.size(1) > num_words:
                          mask = mask[:, :num_words]
  
                      #######################################################
                      # (2) Generate fake images
                      ######################################################
                      noise.data.normal_(0, 1)
                      fake_imgs, _, _, _ = netG(noise, sent_emb, words_embs, mask)
                      for j in range(batch_size):
                          s_tmp = '%s/single/%s' % (save_dir, keys[j])
                          folder = s_tmp[:s_tmp.rfind('/')]
                          if not os.path.isdir(folder):
                              print('Make a new folder: ', folder)
                              mkdir_p(folder)
                          k = -1
                          # for k in range(len(fake_imgs)):
                          im = fake_imgs[k][j].data.cpu().numpy()
                          # [-1, 1] --> [0, 255]
                          im = (im + 1.0) * 127.5
                          im = im.astype(np.uint8)
                          im = np.transpose(im, (1, 2, 0))
                          im = Image.fromarray(im)
                          fullpath = '%s_s%d.png' % (s_tmp, k)
                          im.save(fullpath)
  
      def gen_example(self, data_dic):
          if cfg.TRAIN.NET_G == '':
              print('Error: the path for morels is not found!')
          else:
              # Build and load the generator
              text_encoder = \
                  RNN_ENCODER(self.n_words, nhidden=cfg.TEXT.EMBEDDING_DIM)
              state_dict = \
                  torch.load(cfg.TRAIN.NET_E, map_location=lambda storage, loc: storage)
              text_encoder.load_state_dict(state_dict)
              print('Load text encoder from:', cfg.TRAIN.NET_E)
              text_encoder = text_encoder.cuda()
              text_encoder.eval()
  
              # the path to save generated images
              if cfg.GAN.B_DCGAN:
                  netG = G_DCGAN()
              else:
                  netG = G_NET()
              s_tmp = cfg.TRAIN.NET_G[:cfg.TRAIN.NET_G.rfind('.pth')]
              model_dir = cfg.TRAIN.NET_G
              state_dict = \
                  torch.load(model_dir, map_location=lambda storage, loc: storage)
              netG.load_state_dict(state_dict)
              print('Load G from: ', model_dir)
              netG.cuda()
              netG.eval()
              for key in data_dic:
                  save_dir = '%s/%s' % (s_tmp, key)
                  mkdir_p(save_dir)
                  captions, cap_lens, sorted_indices = data_dic[key]
  
                  batch_size = captions.shape[0]
                  nz = cfg.GAN.Z_DIM
                  captions = Variable(torch.from_numpy(captions), volatile=True)
                  cap_lens = Variable(torch.from_numpy(cap_lens), volatile=True)
  
                  captions = captions.cuda()
                  cap_lens = cap_lens.cuda()
                  for i in range(1):  # 16
                      noise = Variable(torch.FloatTensor(batch_size, nz), volatile=True)
                      noise = noise.cuda()
                      #######################################################
                      # (1) Extract text embeddings
                      ######################################################
                      hidden = text_encoder.init_hidden(batch_size)
                      # words_embs: batch_size x nef x seq_len
                      # sent_emb: batch_size x nef
                      words_embs, sent_emb = text_encoder(captions, cap_lens, hidden)
                      mask = (captions == 0)
                      #######################################################
                      # (2) Generate fake images
                      ######################################################
                      noise.data.normal_(0, 1)
                      fake_imgs, attention_maps, _, _ = netG(noise, sent_emb, words_embs, mask)
                      # G attention
                      cap_lens_np = cap_lens.cpu().data.numpy()
                      for j in range(batch_size):
                          save_name = '%s/%d_s_%d' % (save_dir, i, sorted_indices[j])
                          for k in range(len(fake_imgs)):
                              im = fake_imgs[k][j].data.cpu().numpy()
                              im = (im + 1.0) * 127.5
                              im = im.astype(np.uint8)
                              # print('im', im.shape)
                              im = np.transpose(im, (1, 2, 0))
                              # print('im', im.shape)
                              im = Image.fromarray(im)
                              fullpath = '%s_g%d.png' % (save_name, k)
                              im.save(fullpath)
  
                          for k in range(len(attention_maps)):
                              if len(fake_imgs) > 1:
                                  im = fake_imgs[k + 1].detach().cpu()
                              else:
                                  im = fake_imgs[0].detach().cpu()
                              attn_maps = attention_maps[k]
                              att_sze = attn_maps.size(2)
                              img_set, sentences = \
                                  build_super_images2(im[j].unsqueeze(0),
                                                      captions[j].unsqueeze(0),
                                                      [cap_lens_np[j]], self.ixtoword,
                                                      [attn_maps[j]], att_sze)
                              if img_set is not None:
                                  im = Image.fromarray(img_set)
                                  fullpath = '%s_a%d.png' % (save_name, k)
                                  im.save(fullpath)
  