text: '''
  Python 3.6 
  Pytorch >= 0.4
  Written by Hongyu Wang in Beihang university
  '''
  import torch
  import torch.nn as nn
  from torch.autograd import Variable
  import torch.nn.functional as F
  
  class AttnDecoderRNN(nn.Module):
      def __init__(self, hidden_size, output_size, dropout_p=0.1,):
          super(AttnDecoderRNN, self).__init__()
          self.hidden_size = hidden_size
          self.output_size = output_size
          self.dropout_p = dropout_p
          self.dropout = nn.Dropout(self.dropout_p)
  
          self.embedding = nn.Embedding(self.output_size, 256)
          #self.gru = nn.GRUCell(684, 256)
          self.gru = nn.GRUCell(1024, self.hidden_size)
          self.gru1 = nn.GRUCell(256, self.hidden_size)
          self.out = nn.Linear(128, self.output_size)
          self.hidden = nn.Linear(self.hidden_size, 256)
          self.emb = nn.Linear(256, 128)
          self.conv1 = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=1)
          self.conv_et = nn.Conv2d(1,1,kernel_size=3,stride=1,padding=1)
          self.conv_tan = nn.Conv2d(256,256,kernel_size=3,stride=1,padding=1)
          self.hidden2 = nn.Linear(self.hidden_size, 128)
          self.emb2 = nn.Linear(256, 128)
          self.ua = nn.Linear(1024, 256)
          self.uf = nn.Linear(1, 256)
          self.v = nn.Linear(256, 1)
          self.wc = nn.Linear(1024, 128)
          self.bn = nn.BatchNorm2d(1)
          self.bn1 = nn.BatchNorm2d(256)
          self.relu = nn.ReLU()
  
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.xavier_uniform_(m.weight)
              elif isinstance(m, nn.Linear):
                  nn.init.xavier_uniform_(m.weight)
  
      def forward(self, input_a, hidden, encoder_outputs,bb,attention_sum,decoder_attention,dense_input,batch_size,h_mask,w_mask,gpu):
  
          # batch_gpu must be an int object
          batch_gpu = int(batch_size/len(gpu))
          et_mask = torch.zeros(batch_gpu,dense_input,bb).cuda()
  
          if et_mask.device == torch.device('cuda:0'):
              for i in range(batch_gpu):
                  et_mask[i][:h_mask[i],:w_mask[i]]=1
  
          if et_mask.device == torch.device('cuda:1'):
              for i in range(batch_gpu):
                  et_mask[i][:h_mask[i+1*batch_gpu],:w_mask[i+1*batch_gpu]]=1
  
          if et_mask.device == torch.device('cuda:2'):
              for i in range(batch_gpu):
                  et_mask[i][:h_mask[i+2*batch_gpu],:w_mask[i+2*batch_gpu]]=1
  
          if et_mask.device == torch.device('cuda:3'):
              for i in range(batch_gpu):
                  et_mask[i][:h_mask[i+3*batch_gpu],:w_mask[i+3*batch_gpu]]=1
  
          et_mask_4 = et_mask.unsqueeze(1)
  
          # embedding the word from 1 to 256(total 112 words)
          embedded = self.embedding(input_a).view(batch_gpu,256)
          embedded = self.dropout(embedded)
          hidden = hidden.view(batch_gpu,self.hidden_size)
  
          st = self.gru1(embedded,hidden)
          hidden1 = self.hidden(st)
          hidden1 = hidden1.view(batch_gpu,1,1,256)
  
          # encoder_outputs from (batch,1024,height,width) => (batch,height,width,1024)
          encoder_outputs_trans = torch.transpose(encoder_outputs,1,2)
          encoder_outputs_trans = torch.transpose(encoder_outputs_trans,2,3)
  
          # encoder_outputs_trans (batch,height,width,1024) attention_sum_trans (batch,height,width,1) hidden1 (batch,1,1,256)
          decoder_attention = self.conv1(decoder_attention)
          attention_sum = attention_sum + decoder_attention
          attention_sum_trans = torch.transpose(attention_sum,1,2)
          attention_sum_trans = torch.transpose(attention_sum_trans,2,3)
  
          # encoder_outputs1 (batch,height,width,256) attention_sum1 (batch,height,width,256)
          encoder_outputs1 = self.ua(encoder_outputs_trans)
          attention_sum1 = self.uf(attention_sum_trans)
  
          et = hidden1 + encoder_outputs1 + attention_sum1
          et_trans = torch.transpose(et,2,3)
          et_trans = torch.transpose(et_trans,1,2)
          et_trans = self.conv_tan(et_trans)
          et_trans = et_trans*et_mask_4
          et_trans = self.bn1(et_trans)
          et_trans = torch.tanh(et_trans)
          et_trans = torch.transpose(et_trans,1,2)
          et_trans = torch.transpose(et_trans,2,3)
  
          et = self.v(et_trans) #4,9,34,1
          et = et.squeeze(3)
          # et = torch.transpose(et,2,3)
          # et = torch.transpose(et,1,2)
          # et = self.conv_et(et)
          # et = et*et_mask_4
          # et = self.bn(et)
          # et = self.relu(et) 
          # et = et.squeeze(1)
  
          # et_div_all is attention alpha
          et_div_all = torch.zeros(batch_gpu,1,dense_input,bb)
          et_div_all = et_div_all.cuda()
  
          et_exp = torch.exp(et)
          et_exp = et_exp*et_mask
          et_sum = torch.sum(et_exp,dim=1)
          et_sum = torch.sum(et_sum,dim=1)
          for i in range(batch_gpu):
              et_div = et_exp[i]/(et_sum[i]+1e-8)
              et_div = et_div.unsqueeze(0)
              et_div_all[i] = et_div
  
          # ct is context vector (batch,128)
          ct = et_div_all*encoder_outputs
          ct = ct.sum(dim=2)
          ct = ct.sum(dim=2)
  
          # the next hidden after gru
          # batch,hidden_size
          hidden_next_a = self.gru(ct,st)
          hidden_next = hidden_next_a.view(batch_gpu, 1, self.hidden_size)
  
          # compute the output (batch,128)
          hidden2 = self.hidden2(hidden_next_a)
          embedded2 = self.emb2(embedded)
          ct2 = self.wc(ct)
  
          #output
          # output = F.log_softmax(self.out(hidden2+embedded2+ct2), dim=1)
          out_emb = self.dropout(hidden2+embedded2+ct2)
          output = F.log_softmax(self.out(out_emb), dim=1)
          output = output.unsqueeze(1)
  
          return output, hidden_next, et_div_all, attention_sum, out_emb
  
      def initHidden(self,batch_size):
          result = Variable(torch.randn(batch_size, 1, self.hidden_size))
          return result.cuda()
  