code: import torch
  import torch.nn as nn
  import torch.nn.parallel
  from torch.autograd import Variable
  from torchvision import models
  import torch.utils.model_zoo as model_zoo
  import torch.nn.functional as F
  
  from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
  
  from miscc.config import cfg
  from GlobalAttention import GlobalAttentionGeneral as ATT_NET
  
  # FIXME change this to mutable variables.
  MAX_OBJECTS = cfg.MAX_OBJECTS
  
  def Level2RNNEncodeMagic(captions, cap_lens, low_level_rnn, high_level_rnn):
      ''' You are not expected to understand this. '''
      cap_mask = (cap_lens > 0)
      actual_captions = captions[cap_mask] # -> (? x n_words)
      actual_cap_lens = cap_lens[cap_mask] # -> (?, )
      L_hidden = low_level_rnn.init_hidden(actual_captions.size(0))
      _, actual_cap_embs = low_level_rnn(actual_captions, actual_cap_lens, L_hidden)
      cap_embs = torch.zeros(captions.size(0), cfg.TEXT.MAX_QA_NUM, actual_cap_embs.size(-1)).cuda()
      cap_embs[cap_mask] = actual_cap_embs
      num_caps = cap_mask.sum(1)
      H_hidden = high_level_rnn.init_hidden(captions.size(0))
      per_cap_embs, avg_cap_embs = high_level_rnn(cap_embs, num_caps, H_hidden)
      return per_cap_embs, avg_cap_embs, num_caps
  
  def stn(image, transformation_matrix, size):
      grid = torch.nn.functional.affine_grid(transformation_matrix, torch.Size(size))
      out_image = torch.nn.functional.grid_sample(image, grid)
      return out_image
  
  
  class GLU(nn.Module):
      def __init__(self):
          super(GLU, self).__init__()
  
      def forward(self, x):
          nc = x.size(1)
          assert nc % 2 == 0, 'channels dont divide 2!'
          nc = int(nc/2)
          return x[:, :nc] * torch.sigmoid(x[:, nc:])
  
  
  def conv1x1(in_planes, out_planes, bias=False):
      "1x1 convolution with padding"
      return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1,
                       padding=0, bias=bias)
  
  
  def conv3x3(in_planes, out_planes, stride=1):
      "3x3 convolution with padding"
      return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                       padding=1, bias=False)
  
  
  # Upsale the spatial size by a factor of 2
  def upBlock(in_planes, out_planes):
      block = nn.Sequential(
          # nn.functional.interpolate(scale_factor=2, mode="nearest"),
          nn.Upsample(scale_factor=2, mode='nearest'),
          conv3x3(in_planes, out_planes * 2),
          nn.BatchNorm2d(out_planes * 2),
          GLU())
      return block
  
  
  # Keep the spatial size
  def Block3x3_relu(in_planes, out_planes):
      block = nn.Sequential(
          conv3x3(in_planes, out_planes * 2),
          nn.BatchNorm2d(out_planes * 2),
          GLU())
      return block
  
  
  class ResBlock(nn.Module):
      def __init__(self, channel_num):
          super(ResBlock, self).__init__()
          self.block = nn.Sequential(
              conv3x3(channel_num, channel_num * 2),
              nn.BatchNorm2d(channel_num * 2),
              GLU(),
              conv3x3(channel_num, channel_num),
              nn.BatchNorm2d(channel_num))
  
      def forward(self, x):
          residual = x
          out = self.block(x)
          out += residual
          return out
  
  
  class BBOX_NET(nn.Module):
      # some code is modified from vae examples
      # (https://github.com/pytorch/examples/blob/master/vae/main.py)
      def __init__(self):
          super(BBOX_NET, self).__init__()
          self.c_dim = cfg.GAN.CONDITION_DIM
          self.encode = nn.Sequential(
              # 128 * 16 x 16
              conv3x3(self.c_dim, self.c_dim // 2, stride=2),
              nn.LeakyReLU(0.2, inplace=True),
              # 64 x 8 x 8
              conv3x3(self.c_dim // 2, self.c_dim // 4, stride=2),
              nn.BatchNorm2d(self.c_dim // 4),
              nn.LeakyReLU(0.2, inplace=True),
              # 32 x 4 x 4
              conv3x3(self.c_dim // 4, self.c_dim // 8, stride=2),
              nn.BatchNorm2d(self.c_dim // 8),
              nn.LeakyReLU(0.2, inplace=True),
              # 16 x 2 x 2
          )
  
      def forward(self, labels, transf_matr_inv):
          label_layout = torch.cuda.FloatTensor(labels.shape[0], self.c_dim, 16, 16).fill_(0)
          for idx in range(MAX_OBJECTS):
              current_label = labels[:, idx]
              current_label = current_label.view(current_label.shape[0], current_label.shape[1], 1, 1)
              current_label = current_label.repeat(1, 1, 16, 16)
              current_label = stn(current_label, transf_matr_inv[:, idx], current_label.shape)
              label_layout += current_label
  
          layout_encoding = self.encode(label_layout).view(labels.shape[0], -1)
  
          return layout_encoding
  
  
  # ############## Text2Image Encoder-Decoder #######
  class HigherLevelRNN(nn.Module):
      def __init__(self, ninput=256, drop_prob=0.5,
                   nhidden=128, nlayers=1, bidirectional=True):
          super(HigherLevelRNN, self).__init__()
          self.n_steps = cfg.TEXT.MAX_QA_NUM
          self.ninput = ninput  # size of each embedding vector
          self.drop_prob = drop_prob  # probability of an element to be zeroed
          self.nlayers = nlayers  # Number of recurrent layers
          self.bidirectional = bidirectional
          self.rnn_type = cfg.RNN_TYPE
          if bidirectional:
              self.num_directions = 2
          else:
              self.num_directions = 1
          # number of features in the hidden state
          self.nhidden = nhidden // self.num_directions
  
          self.define_module()
  
      def define_module(self):
          if self.rnn_type == 'LSTM':
              # dropout: If non-zero, introduces a dropout layer on
              # the outputs of each RNN layer except the last layer
              self.rnn = nn.LSTM(self.ninput, self.nhidden,
                                 self.nlayers, batch_first=True,
                                 dropout=self.drop_prob,
                                 bidirectional=self.bidirectional)
          elif self.rnn_type == 'GRU':
              self.rnn = nn.GRU(self.ninput, self.nhidden,
                                self.nlayers, batch_first=True,
                                dropout=self.drop_prob,
                                bidirectional=self.bidirectional)
          else:
              raise NotImplementedError
  
      def init_weights(self):
          return
  
      def init_hidden(self, bsz):
          weight = next(self.parameters()).data
          if self.rnn_type == 'LSTM':
              return (Variable(weight.new(self.nlayers * self.num_directions,
                                          bsz, self.nhidden).zero_()),
                      Variable(weight.new(self.nlayers * self.num_directions,
                                          bsz, self.nhidden).zero_()))
          else:
              return Variable(weight.new(self.nlayers * self.num_directions,
                                         bsz, self.nhidden).zero_())
  
      def forward(self, emb, cap_lens, hidden, mask=None):
          # input: torch.LongTensor of size batch x n_steps
          # emb: batch x n_steps x ninput
          # Returns: a PackedSequence object
          cap_lens = cap_lens.data.tolist()
          emb = pack_padded_sequence(emb, cap_lens, batch_first=True)
          # #hidden and memory (num_layers * num_directions, batch, hidden_size):
          # tensor containing the initial hidden state for each element in batch.
          # #output (batch, seq_len, hidden_size * num_directions)
          # #or a PackedSequence object:
          # tensor containing output features (h_t) from the last layer of RNN
          output, hidden = self.rnn(emb, hidden)
          # PackedSequence object
          # --> (batch, seq_len, hidden_size * num_directions)
          output = pad_packed_sequence(output, batch_first=True)[0]
          # output = self.drop(output)
          # --> batch x hidden_size*num_directions x seq_len
          words_emb = output.transpose(1, 2)
          # --> batch x num_directions*hidden_size
          if self.rnn_type == 'LSTM':
              sent_emb = hidden[0].transpose(0, 1).contiguous()
          else:
              sent_emb = hidden.transpose(0, 1).contiguous()
          sent_emb = sent_emb.view(-1, self.nhidden * self.num_directions)
          return words_emb, sent_emb
  
  class RNN_ENCODER(nn.Module):
      def __init__(self, ntoken, ninput=300, drop_prob=0.5,
                   nhidden=128, nlayers=1, bidirectional=True):
          super(RNN_ENCODER, self).__init__()
          self.n_steps = cfg.TEXT.WORDS_NUM
          self.ntoken = ntoken  # size of the dictionary
          self.ninput = ninput  # size of each embedding vector
          self.drop_prob = drop_prob  # probability of an element to be zeroed
          self.nlayers = nlayers  # Number of recurrent layers
          self.bidirectional = bidirectional
          self.rnn_type = cfg.RNN_TYPE
          if bidirectional:
              self.num_directions = 2
          else:
              self.num_directions = 1
          # number of features in the hidden state
          self.nhidden = nhidden // self.num_directions
  
          self.define_module()
          self.init_weights()
  
      def define_module(self):
          self.encoder = nn.Embedding(self.ntoken, self.ninput)
          self.drop = nn.Dropout(self.drop_prob)
          if self.rnn_type == 'LSTM':
              # dropout: If non-zero, introduces a dropout layer on
              # the outputs of each RNN layer except the last layer
              self.rnn = nn.LSTM(self.ninput, self.nhidden,
                                 self.nlayers, batch_first=True,
                                 dropout=self.drop_prob,
                                 bidirectional=self.bidirectional)
          elif self.rnn_type == 'GRU':
              self.rnn = nn.GRU(self.ninput, self.nhidden,
                                self.nlayers, batch_first=True,
                                dropout=self.drop_prob,
                                bidirectional=self.bidirectional)
          else:
              raise NotImplementedError
  
      def init_weights(self):
          initrange = 0.1
          self.encoder.weight.data.uniform_(-initrange, initrange)
          # Do not need to initialize RNN parameters, which have been initialized
          # http://pytorch.org/docs/master/_modules/torch/nn/modules/rnn.html#LSTM
          # self.decoder.weight.data.uniform_(-initrange, initrange)
          # self.decoder.bias.data.fill_(0)
  
      def init_hidden(self, bsz):
          weight = next(self.parameters()).data
          if self.rnn_type == 'LSTM':
              return (Variable(weight.new(self.nlayers * self.num_directions,
                                          bsz, self.nhidden).zero_()),
                      Variable(weight.new(self.nlayers * self.num_directions,
                                          bsz, self.nhidden).zero_()))
          else:
              return Variable(weight.new(self.nlayers * self.num_directions,
                                         bsz, self.nhidden).zero_())
  
      def forward(self, captions, cap_lens, hidden, mask=None):
          # input: torch.LongTensor of size batch x n_steps
          # --> emb: batch x n_steps x ninput
          emb = self.drop(self.encoder(captions))
          #
          # Returns: a PackedSequence object
          cap_lens = cap_lens.data.tolist()
          emb = pack_padded_sequence(emb, cap_lens, batch_first=True, enforce_sorted=False)
          # #hidden and memory (num_layers * num_directions, batch, hidden_size):
          # tensor containing the initial hidden state for each element in batch.
          # #output (batch, seq_len, hidden_size * num_directions)
          # #or a PackedSequence object:
          # tensor containing output features (h_t) from the last layer of RNN
          output, hidden = self.rnn(emb, hidden)
          # PackedSequence object
          # --> (batch, seq_len, hidden_size * num_directions)
          output = pad_packed_sequence(output, batch_first=True)[0]
          # output = self.drop(output)
          # --> batch x hidden_size*num_directions x seq_len
          words_emb = output.transpose(1, 2)
          # --> batch x num_directions*hidden_size
          if self.rnn_type == 'LSTM':
              sent_emb = hidden[0].transpose(0, 1).contiguous()
          else:
              sent_emb = hidden.transpose(0, 1).contiguous()
          sent_emb = sent_emb.view(-1, self.nhidden * self.num_directions)
          return words_emb, sent_emb
  
  
  class CNN_ENCODER(nn.Module):
      def __init__(self, nef):
          super(CNN_ENCODER, self).__init__()
          if cfg.TRAIN.FLAG:
              self.nef = nef
          else:
              self.nef = 256  # define a uniform ranker
  
          model = models.inception_v3()
          url = 'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth'
          model.load_state_dict(model_zoo.load_url(url))
          for param in model.parameters():
              param.requires_grad = False
          print('Load pretrained model from ', url)
          # print(model)
  
          self.define_module(model)
          self.init_trainable_weights()
  
      def define_module(self, model):
          self.Conv2d_1a_3x3 = model.Conv2d_1a_3x3
          self.Conv2d_2a_3x3 = model.Conv2d_2a_3x3
          self.Conv2d_2b_3x3 = model.Conv2d_2b_3x3
          self.Conv2d_3b_1x1 = model.Conv2d_3b_1x1
          self.Conv2d_4a_3x3 = model.Conv2d_4a_3x3
          self.Mixed_5b = model.Mixed_5b
          self.Mixed_5c = model.Mixed_5c
          self.Mixed_5d = model.Mixed_5d
          self.Mixed_6a = model.Mixed_6a
          self.Mixed_6b = model.Mixed_6b
          self.Mixed_6c = model.Mixed_6c
          self.Mixed_6d = model.Mixed_6d
          self.Mixed_6e = model.Mixed_6e
          self.Mixed_7a = model.Mixed_7a
          self.Mixed_7b = model.Mixed_7b
          self.Mixed_7c = model.Mixed_7c
  
          self.emb_features = conv1x1(768, self.nef)
          self.emb_cnn_code = nn.Linear(2048, self.nef)
  
      def init_trainable_weights(self):
          initrange = 0.1
          self.emb_features.weight.data.uniform_(-initrange, initrange)
          self.emb_cnn_code.weight.data.uniform_(-initrange, initrange)
  
      def forward(self, x):
          features = None
          # --> fixed-size input: batch x 3 x 299 x 299
          # x = nn.functional.interpolate(x, size=(299, 299), mode='bilinear')
          x = nn.Upsample(size=(299, 299), mode='bilinear')(x)
          # 299 x 299 x 3
          x = self.Conv2d_1a_3x3(x)
          # 149 x 149 x 32
          x = self.Conv2d_2a_3x3(x)
          # 147 x 147 x 32
          x = self.Conv2d_2b_3x3(x)
          # 147 x 147 x 64
          x = F.max_pool2d(x, kernel_size=3, stride=2)
          # 73 x 73 x 64
          x = self.Conv2d_3b_1x1(x)
          # 73 x 73 x 80
          x = self.Conv2d_4a_3x3(x)
          # 71 x 71 x 192
  
          x = F.max_pool2d(x, kernel_size=3, stride=2)
          # 35 x 35 x 192
          x = self.Mixed_5b(x)
          # 35 x 35 x 256
          x = self.Mixed_5c(x)
          # 35 x 35 x 288
          x = self.Mixed_5d(x)
          # 35 x 35 x 288
  
          x = self.Mixed_6a(x)
          # 17 x 17 x 768
          x = self.Mixed_6b(x)
          # 17 x 17 x 768
          x = self.Mixed_6c(x)
          # 17 x 17 x 768
          x = self.Mixed_6d(x)
          # 17 x 17 x 768
          x = self.Mixed_6e(x)
          # 17 x 17 x 768
  
          # image region features
          features = x
          # 17 x 17 x 768
  
          x = self.Mixed_7a(x)
          # 8 x 8 x 1280
          x = self.Mixed_7b(x)
          # 8 x 8 x 2048
          x = self.Mixed_7c(x)
          # 8 x 8 x 2048
          x = F.avg_pool2d(x, kernel_size=8)
          # 1 x 1 x 2048
          # x = F.dropout(x, training=self.training)
          # 1 x 1 x 2048
          x = x.view(x.size(0), -1)
          # 2048
  
          # global image features
          cnn_code = self.emb_cnn_code(x)
          # 512
          if features is not None:
              features = self.emb_features(features)
          return features, cnn_code
  
  
  # ############## G networks ###################
  class SimpleLrPhiProjector(nn.Module):
      ''' a fake CA_NET. '''
      def __init__(self):
          super(SimpleLrPhiProjector, self).__init__()
          self.t_dim = cfg.TEXT.EMBEDDING_DIM
          self.c_dim = cfg.GAN.CONDITION_DIM
          self.fc = nn.Linear(self.t_dim, self.c_dim, bias=True)
          self.relu = nn.ReLU()
      def forward(self, x):
          return self.relu(self.fc(x)), None, None
  
  
  class CA_NET(nn.Module):
      # some code is modified from vae examples
      # (https://github.com/pytorch/examples/blob/master/vae/main.py)
      def __init__(self):
          super(CA_NET, self).__init__()
          self.t_dim = cfg.TEXT.EMBEDDING_DIM
          self.c_dim = cfg.GAN.CONDITION_DIM
          self.fc = nn.Linear(self.t_dim, self.c_dim * 4, bias=True)
          self.relu = GLU()
  
      def encode(self, text_embedding):
          x = self.relu(self.fc(text_embedding))
          mu = x[:, :self.c_dim]
          logvar = x[:, self.c_dim:]
          return mu, logvar
  
      def reparametrize(self, mu, logvar):
          std = logvar.mul(0.5).exp_()
          if cfg.CUDA:
              eps = torch.cuda.FloatTensor(std.size()).normal_()
          else:
              eps = torch.FloatTensor(std.size()).normal_()
          eps = Variable(eps)
          return eps.mul(std).add_(mu)
  
      def forward(self, text_embedding):
          mu, logvar = self.encode(text_embedding)
          c_code = self.reparametrize(mu, logvar)
          return c_code, mu, logvar
  
  
  class INIT_STAGE_G(nn.Module):
      def __init__(self, ngf, ncf):
          super(INIT_STAGE_G, self).__init__()
          self.gf_dim = ngf
          self.in_dim = cfg.GAN.Z_DIM + ncf  # cfg.TEXT.EMBEDDING_DIM
  
          self.define_module()
  
      def define_module(self):
          nz, ngf = self.in_dim, self.gf_dim
          linput = 100+cfg.GAN.LABEL_DIM
          self.ef_dim = 100
  
          self.bbox_net = BBOX_NET()
          nz += 48
  
          self.fc = nn.Sequential(
              nn.Linear(nz, ngf * 4 * 4 * 2, bias=False),
              nn.BatchNorm1d(ngf * 4 * 4 * 2),
              GLU())
  
          # local pathway
          self.label = nn.Sequential(
              nn.Linear(linput, self.ef_dim, bias=False),
              nn.BatchNorm1d(self.ef_dim),
              nn.ReLU(True))
          self.local1 = upBlock(self.ef_dim, ngf // 2)
          self.local2 = upBlock(ngf // 2, ngf // 4)
  
          self.upsample1 = upBlock(ngf, ngf // 2)
          self.upsample2 = upBlock(ngf // 2, ngf // 4)
          self.upsample3 = upBlock(ngf // 2, ngf // 8)
          self.upsample4 = upBlock(ngf // 8, ngf // 16)
  
      def forward(self, z_code, c_code, transf_matrices_inv, label_one_hot):
          """
          :param z_code: batch x cfg.GAN.Z_DIM
          :param c_code: batch x cfg.TEXT.EMBEDDING_DIM
          :return: batch x ngf/16 x 64 x 64
          """
          local_labels = torch.cuda.FloatTensor(z_code.shape[0], MAX_OBJECTS, self.ef_dim).fill_(0)
  
          # local pathway
          h_code_locals = torch.cuda.FloatTensor(z_code.shape[0], self.gf_dim // 4, 16, 16).fill_(0)
  
          for idx in range(MAX_OBJECTS):
              current_label = self.label(torch.cat((c_code, label_one_hot[:, idx]), 1))
              local_labels[:, idx] = current_label
              current_label = current_label.view(current_label.shape[0], self.ef_dim, 1, 1)
              current_label = current_label.repeat(1, 1, 4, 4)
              h_code_local = self.local1(current_label)
              h_code_local = self.local2(h_code_local)
              h_code_local = stn(h_code_local, transf_matrices_inv[:, idx], h_code_local.shape)
              h_code_locals += h_code_local
  
          bbox_code = self.bbox_net(local_labels, transf_matrices_inv)
          c_z_code = torch.cat((c_code, z_code, bbox_code), 1)
          # c_z_code = torch.cat((c_code, z_code), 1)
          # state size ngf x 4 x 4
          out_code = self.fc(c_z_code)
          out_code = out_code.view(-1, self.gf_dim, 4, 4)
          # state size ngf/3 x 8 x 8
          out_code = self.upsample1(out_code)
          # state size ngf/4 x 16 x 16
          out_code = self.upsample2(out_code)
  
          # combine local and global pathways
          out_code = torch.cat((out_code, h_code_locals), 1)
  
          # state size ngf/8 x 32 x 32
          out_code32 = self.upsample3(out_code)
          # state size ngf/16 x 64 x 64
          out_code64 = self.upsample4(out_code32)
  
          return out_code64
  
  
  class NEXT_STAGE_G(nn.Module):
      def __init__(self, ngf, nef, ncf):
          super(NEXT_STAGE_G, self).__init__()
          self.gf_dim = ngf
          self.ef_dim = nef
          self.cf_dim = ncf
          self.num_residual = cfg.GAN.R_NUM
          self.define_module()
  
      def _make_layer(self, block, channel_num):
          layers = []
          for i in range(cfg.GAN.R_NUM):
              layers.append(block(channel_num))
          return nn.Sequential(*layers)
  
      def define_module(self):
          ngf = self.gf_dim
          self.att = ATT_NET(ngf, self.ef_dim)
          self.jointConv = Block3x3_relu(ngf + 100, ngf) # FIXME del
          self.residual = self._make_layer(ResBlock, ngf) # FIXME ngf * 2
          # self.upsample = upBlock(ngf * 2, ngf) # FIXME
          self.upsample = upBlock(ngf , ngf)
  
      def forward(self, h_code, c_code, word_embs, mask):
          """
              h_code1(query):  batch x idf x ih x iw (queryL=ihxiw)
              word_embs(context): batch x cdf x sourceL (sourceL=seq_len)
              c_code1: batch x idf x queryL
              att1: batch x sourceL x queryL
          """
          self.att.applyMask(mask)
          _c_code, att = self.att(h_code, word_embs)
  
          # FIXME del
          s_size = h_code.size(2)
          c_code = c_code.view(-1, 100, 1, 1)
          c_code = c_code.repeat(1, 1, s_size, s_size)
  
          h_c_code = torch.cat((h_code, c_code), 1)
          h_c_code = self.jointConv(h_c_code) # FIXME del
          out_code = self.residual(h_c_code)
  
          # state size ngf/2 x 2in_size x 2in_size
          out_code = self.upsample(out_code)
  
          return out_code, att
  
  
  class GET_IMAGE_G(nn.Module):
      def __init__(self, ngf):
          super(GET_IMAGE_G, self).__init__()
          self.gf_dim = ngf
          self.img = nn.Sequential(
              conv3x3(ngf, 3),
              nn.Tanh()
          )
  
      def forward(self, h_code):
          out_img = self.img(h_code)
          return out_img
  
  
  class G_NET(nn.Module):
      def __init__(self):
          super(G_NET, self).__init__()
          ngf = cfg.GAN.GF_DIM
          nef = cfg.TEXT.EMBEDDING_DIM
          ncf = cfg.GAN.CONDITION_DIM
          
          if cfg.GAN.B_CA_NET:
              self.ca_net = CA_NET()
          else:
              self.ca_net = SimpleLrPhiProjector()
  
          if cfg.TREE.BRANCH_NUM > 0:
              self.h_net1 = INIT_STAGE_G(ngf * 16, ncf)
              self.img_net1 = GET_IMAGE_G(ngf)
          # gf x 64 x 64
          if cfg.TREE.BRANCH_NUM > 1:
              self.h_net2 = NEXT_STAGE_G(ngf, nef, ncf)
              self.img_net2 = GET_IMAGE_G(ngf)
          if cfg.TREE.BRANCH_NUM > 2:
              self.h_net3 = NEXT_STAGE_G(ngf, nef, ncf)
              self.img_net3 = GET_IMAGE_G(ngf)
  
      def forward(self, z_code, sent_emb, word_embs, mask, transf_matrices_inv, label_one_hot):
          """
              :param z_code: batch x cfg.GAN.Z_DIM
              :param sent_emb: batch x cfg.TEXT.EMBEDDING_DIM
              :param word_embs: batch x cdf x seq_len
              :param mask: batch x seq_len
              :return:
          """
          fake_imgs = []
          att_maps = []
          c_code, mu, logvar = self.ca_net(sent_emb)
  
          if cfg.TREE.BRANCH_NUM > 0:
              h_code1 = self.h_net1(z_code, c_code, transf_matrices_inv, label_one_hot)
              fake_img1 = self.img_net1(h_code1)
              fake_imgs.append(fake_img1)
          if cfg.TREE.BRANCH_NUM > 1:
              h_code2, att1 = \
                  self.h_net2(h_code1, c_code, word_embs, mask)
              fake_img2 = self.img_net2(h_code2)
              fake_imgs.append(fake_img2)
              if att1 is not None:
                  att_maps.append(att1)
          if cfg.TREE.BRANCH_NUM > 2:
              h_code3, att2 = \
                  self.h_net3(h_code2, c_code, word_embs, mask)
              fake_img3 = self.img_net3(h_code3)
              fake_imgs.append(fake_img3)
              if att2 is not None:
                  att_maps.append(att2)
  
          return fake_imgs, att_maps, mu, logvar
  
  
  class G_DCGAN(nn.Module):
      def __init__(self):
          super(G_DCGAN, self).__init__()
          ngf = cfg.GAN.GF_DIM
          nef = cfg.TEXT.EMBEDDING_DIM
          ncf = cfg.GAN.CONDITION_DIM
          self.ca_net = CA_NET()
  
          # 16gf x 64 x 64 --> gf x 64 x 64 --> 3 x 64 x 64
          if cfg.TREE.BRANCH_NUM > 0:
              self.h_net1 = INIT_STAGE_G(ngf * 16, ncf)
          # gf x 64 x 64
          if cfg.TREE.BRANCH_NUM > 1:
              self.h_net2 = NEXT_STAGE_G(ngf, nef, ncf)
          if cfg.TREE.BRANCH_NUM > 2:
              self.h_net3 = NEXT_STAGE_G(ngf, nef, ncf)
          self.img_net = GET_IMAGE_G(ngf)
  
      def forward(self, z_code, sent_emb, word_embs, mask):
          """
              :param z_code: batch x cfg.GAN.Z_DIM
              :param sent_emb: batch x cfg.TEXT.EMBEDDING_DIM
              :param word_embs: batch x cdf x seq_len
              :param mask: batch x seq_len
              :return:
          """
          att_maps = []
          c_code, mu, logvar = self.ca_net(sent_emb)
          if cfg.TREE.BRANCH_NUM > 0:
              h_code = self.h_net1(z_code, c_code)
          if cfg.TREE.BRANCH_NUM > 1:
              h_code, att1 = self.h_net2(h_code, c_code, word_embs, mask)
              if att1 is not None:
                  att_maps.append(att1)
          if cfg.TREE.BRANCH_NUM > 2:
              h_code, att2 = self.h_net3(h_code, c_code, word_embs, mask)
              if att2 is not None:
                  att_maps.append(att2)
  
          fake_imgs = self.img_net(h_code)
          return [fake_imgs], att_maps, mu, logvar
  
  
  # ############## D networks ##########################
  def Block3x3_leakRelu(in_planes, out_planes):
      block = nn.Sequential(
          conv3x3(in_planes, out_planes),
          nn.BatchNorm2d(out_planes),
          nn.LeakyReLU(0.2, inplace=True)
      )
      return block
  
  
  # Downsale the spatial size by a factor of 2
  def downBlock(in_planes, out_planes):
      block = nn.Sequential(
          nn.Conv2d(in_planes, out_planes, 4, 2, 1, bias=False),
          nn.BatchNorm2d(out_planes),
          nn.LeakyReLU(0.2, inplace=True)
      )
      return block
  
  
  # Downsale the spatial size by a factor of 16
  def encode_image_by_16times(ndf):
      encode_img = nn.Sequential(
          # --> state size. ndf x in_size/2 x in_size/2
          nn.Conv2d(3, ndf, 4, 2, 1, bias=False),
          nn.LeakyReLU(0.2, inplace=True),
          # --> state size 2ndf x x in_size/4 x in_size/4
          nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
          nn.BatchNorm2d(ndf * 2),
          nn.LeakyReLU(0.2, inplace=True),
          # --> state size 4ndf x in_size/8 x in_size/8
          nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
          nn.BatchNorm2d(ndf * 4),
          nn.LeakyReLU(0.2, inplace=True),
          # --> state size 8ndf x in_size/16 x in_size/16
          nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
          nn.BatchNorm2d(ndf * 8),
          nn.LeakyReLU(0.2, inplace=True)
      )
      return encode_img
  
  
  class D_GET_LOGITS(nn.Module):
      def __init__(self, ndf, nef, bcondition=False):
          super(D_GET_LOGITS, self).__init__()
          self.df_dim = ndf
          self.ef_dim = nef
          self.bcondition = bcondition
          if self.bcondition:
              self.jointConv = Block3x3_leakRelu(ndf * 8 + nef, ndf * 8)
  
          self.outlogits = nn.Sequential(
              nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=4),
              nn.Sigmoid())
  
      def forward(self, h_code, c_code=None):
          if self.bcondition and c_code is not None:
              # conditioning output
              c_code = c_code.view(-1, self.ef_dim, 1, 1)
              c_code = c_code.repeat(1, 1, 4, 4)
              # state size (ngf+egf) x 4 x 4
              h_c_code = torch.cat((h_code, c_code), 1)
              # state size ngf x in_size x in_size
              h_c_code = self.jointConv(h_c_code)
          else:
              h_c_code = h_code
  
          output = self.outlogits(h_c_code)
          return output.view(-1)
  
  
  # For 64 x 64 images
  class D_NET64(nn.Module):
      def __init__(self, b_jcu=True):
          super(D_NET64, self).__init__()
          ndf = cfg.GAN.DF_DIM
          nef = cfg.TEXT.EMBEDDING_DIM
          if b_jcu:
              self.UNCOND_DNET = D_GET_LOGITS(ndf, nef, bcondition=False)
          else:
              self.UNCOND_DNET = None
          self.COND_DNET = D_GET_LOGITS(ndf, nef, bcondition=True)
          self.define_module()
  
      def define_module(self):
          self.act = nn.LeakyReLU(0.2, inplace=True)
          ndf = cfg.GAN.DF_DIM
  
          # global pathway
          # --> state size. ndf x in_size/2 x in_size/2
          self.conv1 = nn.Conv2d(3, ndf, 4, 2, 1, bias=False)
          # --> state size 2ndf x x in_size/4 x in_size/4
          self.conv2 = nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False)
          self.bn2 = nn.BatchNorm2d(ndf * 2)
          # --> state size 4ndf x in_size/8 x in_size/8
          self.conv3 = nn.Conv2d(ndf * 4, ndf * 4, 4, 2, 1, bias=False)
          self.bn3 = nn.BatchNorm2d(ndf * 4)
          # --> state size 8ndf x in_size/16 x in_size/16
          self.conv4 = nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False)
          self.bn4 = nn.BatchNorm2d(ndf * 8)
  
          # local pathway
          self.local = nn.Sequential(
              nn.Conv2d(3 + cfg.GAN.LABEL_DIM, ndf * 2, 4, 1, 1, bias=False),
              nn.BatchNorm2d(ndf * 2),
              nn.LeakyReLU(0.2, inplace=True)
          )
  
      def forward(self, image, label, transf_matrices, transf_matrices_inv):
          # local pathway
          h_code_locals = torch.cuda.FloatTensor(image.shape[0], cfg.GAN.DF_DIM * 2, 16, 16).fill_(0)
          for idx in range(MAX_OBJECTS):
              current_label = label[:, idx].view(label.shape[0], cfg.GAN.LABEL_DIM, 1, 1)
              current_label = current_label.repeat(1, 1, 16, 16)
              h_code_local = stn(image, transf_matrices[:, idx], (image.shape[0], image.shape[1], 16, 16))
              h_code_local = torch.cat((h_code_local, current_label), 1)
              h_code_local = self.local(h_code_local)
              h_code_local = stn(h_code_local, transf_matrices_inv[:, idx],
                                 (h_code_local.shape[0], h_code_local.shape[1], 16, 16))
              h_code_locals += h_code_local
  
          h_code = self.conv1(image)
          h_code = self.act(h_code)
          h_code = self.conv2(h_code)
          h_code = self.bn2(h_code)
          h_code = self.act(h_code)
  
          h_code = torch.cat((h_code, h_code_locals), 1)
  
          h_code = self.conv3(h_code)
          h_code = self.bn3(h_code)
          h_code = self.act(h_code)
  
          h_code = self.conv4(h_code)
          h_code = self.bn4(h_code)
          x_code4 = self.act(h_code)
  
          return x_code4
  
  
  # For 128 x 128 images
  class D_NET128(nn.Module):
      def __init__(self, b_jcu=True):
          super(D_NET128, self).__init__()
          ndf = cfg.GAN.DF_DIM
          nef = cfg.TEXT.EMBEDDING_DIM
          self.img_code_s16 = encode_image_by_16times(ndf)
          self.img_code_s32 = downBlock(ndf * 8, ndf * 16)
          self.img_code_s32_1 = Block3x3_leakRelu(ndf * 16, ndf * 8)
          #
          if b_jcu:
              self.UNCOND_DNET = D_GET_LOGITS(ndf, nef, bcondition=False)
          else:
              self.UNCOND_DNET = None
          self.COND_DNET = D_GET_LOGITS(ndf, nef, bcondition=True)
  
      def forward(self, x_var):
          x_code8 = self.img_code_s16(x_var)   # 8 x 8 x 8df
          x_code4 = self.img_code_s32(x_code8)   # 4 x 4 x 16df
          x_code4 = self.img_code_s32_1(x_code4)  # 4 x 4 x 8df
          return x_code4
  
  
  # For 256 x 256 images
  class D_NET256(nn.Module):
      def __init__(self, b_jcu=True):
          super(D_NET256, self).__init__()
          ndf = cfg.GAN.DF_DIM
          nef = cfg.TEXT.EMBEDDING_DIM
          self.img_code_s16 = encode_image_by_16times(ndf)
          self.img_code_s32 = downBlock(ndf * 8, ndf * 16)
          self.img_code_s64 = downBlock(ndf * 16, ndf * 32)
          self.img_code_s64_1 = Block3x3_leakRelu(ndf * 32, ndf * 16)
          self.img_code_s64_2 = Block3x3_leakRelu(ndf * 16, ndf * 8)
          if b_jcu:
              self.UNCOND_DNET = D_GET_LOGITS(ndf, nef, bcondition=False)
          else:
              self.UNCOND_DNET = None
          self.COND_DNET = D_GET_LOGITS(ndf, nef, bcondition=True)
  
      def forward(self, x_var):
          x_code16 = self.img_code_s16(x_var)
          x_code8 = self.img_code_s32(x_code16)
          x_code4 = self.img_code_s64(x_code8)
          x_code4 = self.img_code_s64_1(x_code4)
          x_code4 = self.img_code_s64_2(x_code4)
          return x_code4
  
apis: torch.nn.functional.avg_pool2d torch.nn.GRU torch.nn.Upsample torch.nn.Conv2d torch.nn.functional.grid_sample torch.nn.Sigmoid torch.cat torch.nn.utils.rnn.pad_packed_sequence torch.nn.Dropout torch.nn.BatchNorm2d torch.nn.BatchNorm1d torch.autograd.Variable torch.utils.model_zoo.load_url torch.sigmoid torch.nn.LSTM torch.nn.functional.max_pool2d torch.Size torch.nn.utils.rnn.pack_padded_sequence torch.nn.Linear torch.cuda.FloatTensor torch.nn.Embedding torch.nn.Tanh torch.nn.Sequential torch.nn.ReLU torch.nn.LeakyReLU