code: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  import numpy as np
  
  from .layers.contrastive import ContrastiveLoss
  from .layers.utils import l1norm, l2norm
  from .layers.img_enc import EncoderImage
  from .layers.txt_enc import EncoderText
  
  
  class VisualSA(nn.Module):
      """
      Build global image representations by self-attention.
      Args: - local: local region embeddings, shape: (batch_size, 36, 1024)
            - raw_global: raw image by averaging regions, shape: (batch_size, 1024)
      Returns: - new_global: final image by self-attention, shape: (batch_size, 1024).
      """
      def __init__(self, embed_dim, dropout_rate, num_region):
          super(VisualSA, self).__init__()
  
          self.embedding_local = nn.Sequential(nn.Linear(embed_dim, embed_dim),
                                               nn.BatchNorm1d(num_region),
                                               nn.Tanh(), nn.Dropout(dropout_rate))
          self.embedding_global = nn.Sequential(nn.Linear(embed_dim, embed_dim),
                                                nn.BatchNorm1d(embed_dim),
                                                nn.Tanh(), nn.Dropout(dropout_rate))
          self.embedding_common = nn.Sequential(nn.Linear(embed_dim, 1))
  
          self.init_weights()
          self.softmax = nn.Softmax(dim=1)
  
      def init_weights(self):
          for embeddings in self.children():
              for m in embeddings:
                  if isinstance(m, nn.Linear):
                      r = np.sqrt(6.) / np.sqrt(m.in_features + m.out_features)
                      m.weight.data.uniform_(-r, r)
                      m.bias.data.fill_(0)
                  elif isinstance(m, nn.BatchNorm1d):
                      m.weight.data.fill_(1)
                      m.bias.data.zero_()
  
      def forward(self, local, raw_global):
          # compute embedding of local regions and raw global image
          l_emb = self.embedding_local(local)
          g_emb = self.embedding_global(raw_global)
  
          # compute the normalized weights, shape: (batch_size, 36)
          g_emb = g_emb.unsqueeze(1).repeat(1, l_emb.size(1), 1)
          common = l_emb.mul(g_emb)
          weights = self.embedding_common(common).squeeze(2)
          weights = self.softmax(weights)
  
          # compute final image, shape: (batch_size, 1024)
          new_global = (weights.unsqueeze(2) * local).sum(dim=1)
          new_global = l2norm(new_global, dim=-1)
  
          return new_global
  
  
  class TextSA(nn.Module):
      """
      Build global text representations by self-attention.
      Args: - local: local word embeddings, shape: (batch_size, L, 1024)
            - raw_global: raw text by averaging words, shape: (batch_size, 1024)
      Returns: - new_global: final text by self-attention, shape: (batch_size, 1024).
      """
  
      def __init__(self, embed_dim, dropout_rate):
          super(TextSA, self).__init__()
  
          self.embedding_local = nn.Sequential(nn.Linear(embed_dim, embed_dim),
                                               nn.Tanh(), nn.Dropout(dropout_rate))
          self.embedding_global = nn.Sequential(nn.Linear(embed_dim, embed_dim),
                                                nn.Tanh(), nn.Dropout(dropout_rate))
          self.embedding_common = nn.Sequential(nn.Linear(embed_dim, 1))
  
          self.init_weights()
          self.softmax = nn.Softmax(dim=1)
  
      def init_weights(self):
          for embeddings in self.children():
              for m in embeddings:
                  if isinstance(m, nn.Linear):
                      r = np.sqrt(6.) / np.sqrt(m.in_features + m.out_features)
                      m.weight.data.uniform_(-r, r)
                      m.bias.data.fill_(0)
                  elif isinstance(m, nn.BatchNorm1d):
                      m.weight.data.fill_(1)
                      m.bias.data.zero_()
  
      def forward(self, local, raw_global):
          # compute embedding of local words and raw global text
          l_emb = self.embedding_local(local)
          g_emb = self.embedding_global(raw_global)
  
          # compute the normalized weights, shape: (batch_size, L)
          g_emb = g_emb.unsqueeze(1).repeat(1, l_emb.size(1), 1)
          common = l_emb.mul(g_emb)
          weights = self.embedding_common(common).squeeze(2)
          weights = self.softmax(weights)
  
          # compute final text, shape: (batch_size, 1024)
          new_global = (weights.unsqueeze(2) * local).sum(dim=1)
          new_global = l2norm(new_global, dim=-1)
  
          return new_global
  
  
  class GraphReasoning(nn.Module):
      """
      Perform the similarity graph reasoning with a full-connected graph
      Args: - sim_emb: global and local alignments, shape: (batch_size, L+1, 256)
      Returns; - sim_sgr: reasoned graph nodes after several steps, shape: (batch_size, L+1, 256)
      """
      def __init__(self, sim_dim):
          super(GraphReasoning, self).__init__()
  
          self.graph_query_w = nn.Linear(sim_dim, sim_dim)
          self.graph_key_w = nn.Linear(sim_dim, sim_dim)
          self.sim_graph_w = nn.Linear(sim_dim, sim_dim)
          self.relu = nn.ReLU()
  
          self.init_weights()
  
      def forward(self, sim_emb):
          sim_query = self.graph_query_w(sim_emb)
          sim_key = self.graph_key_w(sim_emb)
          sim_edge = torch.softmax(torch.bmm(sim_query, sim_key.permute(0, 2, 1)), dim=-1)
          sim_sgr = torch.bmm(sim_edge, sim_emb)
          sim_sgr = self.relu(self.sim_graph_w(sim_sgr))
          return sim_sgr
  
      def init_weights(self):
          for m in self.children():
              if isinstance(m, nn.Linear):
                  r = np.sqrt(6.) / np.sqrt(m.in_features + m.out_features)
                  m.weight.data.uniform_(-r, r)
                  m.bias.data.fill_(0)
              elif isinstance(m, nn.BatchNorm1d):
                  m.weight.data.fill_(1)
                  m.bias.data.zero_()
  
  
  class AttentionFiltration(nn.Module):
      """
      Perform the similarity Attention Filtration with a gate-based attention
      Args: - sim_emb: global and local alignments, shape: (batch_size, L+1, 256)
      Returns; - sim_saf: aggregated alignment after attention filtration, shape: (batch_size, 256)
      """
      def __init__(self, sim_dim):
          super(AttentionFiltration, self).__init__()
  
          self.attn_sim_w = nn.Linear(sim_dim, 1)
          self.bn = nn.BatchNorm1d(1)
  
          self.init_weights()
  
      def forward(self, sim_emb):
          sim_attn = l1norm(torch.sigmoid(self.bn(self.attn_sim_w(sim_emb).permute(0, 2, 1))), dim=-1)
          sim_saf = torch.matmul(sim_attn, sim_emb)
          sim_saf = l2norm(sim_saf.squeeze(1), dim=-1)
          return sim_saf
  
      def init_weights(self):
          for m in self.children():
              if isinstance(m, nn.Linear):
                  r = np.sqrt(6.) / np.sqrt(m.in_features + m.out_features)
                  m.weight.data.uniform_(-r, r)
                  m.bias.data.fill_(0)
              elif isinstance(m, nn.BatchNorm1d):
                  m.weight.data.fill_(1)
                  m.bias.data.zero_()
  
  
  class EncoderSimilarity(nn.Module):
      """
      Compute the image-text similarity by SGR, SAF, AVE
      Args: - img_emb: local region embeddings, shape: (batch_size, 36, 1024)
            - cap_emb: local word embeddings, shape: (batch_size, L, 1024)
      Returns:
          - sim_all: final image-text similarities, shape: (batch_size, batch_size).
      """
      def __init__(self, embed_size, sim_dim, module_name='AVE', sgr_step=3):
          super(EncoderSimilarity, self).__init__()
          self.module_name = module_name
  
          self.v_global_w = VisualSA(embed_size, 0.4, 36)
          self.t_global_w = TextSA(embed_size, 0.4)
  
          self.sim_tranloc_w = nn.Linear(embed_size, sim_dim)
          self.sim_tranglo_w = nn.Linear(embed_size, sim_dim)
  
          self.sim_eval_w = nn.Linear(sim_dim, 1)
          self.sigmoid = nn.Sigmoid()
  
          if module_name == 'SGR':
              self.SGR_module = nn.ModuleList([GraphReasoning(sim_dim) for i in range(sgr_step)])
          elif module_name == 'SAF':
              self.SAF_module = AttentionFiltration(sim_dim)
          else:
              raise ValueError('Invalid input of opt.module_name in opts.py')
  
          self.init_weights()
  
      def forward(self, img_emb, cap_emb, cap_lens):
          sim_all = []
          n_image = img_emb.size(0)
          n_caption = cap_emb.size(0)
  
          # get enhanced global images by self-attention
          img_ave = torch.mean(img_emb, 1)
          img_glo = self.v_global_w(img_emb, img_ave)
  
          for i in range(n_caption):
              # get the i-th sentence
              n_word = cap_lens[i]
              cap_i = cap_emb[i, :n_word, :].unsqueeze(0)
              cap_i_expand = cap_i.repeat(n_image, 1, 1)
  
              # get enhanced global i-th text by self-attention
              cap_ave_i = torch.mean(cap_i, 1)
              cap_glo_i = self.t_global_w(cap_i, cap_ave_i)
  
              # local-global alignment construction
              Context_img = SCAN_attention(cap_i_expand, img_emb, smooth=9.0)
              sim_loc = torch.pow(torch.sub(Context_img, cap_i_expand), 2)
              sim_loc = l2norm(self.sim_tranloc_w(sim_loc), dim=-1)
  
              sim_glo = torch.pow(torch.sub(img_glo, cap_glo_i), 2)
              sim_glo = l2norm(self.sim_tranglo_w(sim_glo), dim=-1)
  
              # concat the global and local alignments
              sim_emb = torch.cat([sim_glo.unsqueeze(1), sim_loc], 1)
  
              # compute the final similarity vector
              if self.module_name == 'SGR':
                  for module in self.SGR_module:
                      sim_emb = module(sim_emb)
                  sim_vec = sim_emb[:, 0, :]
              else:
                  sim_vec = self.SAF_module(sim_emb)
  
              # compute the final similarity score
              sim_i = self.sigmoid(self.sim_eval_w(sim_vec))
              sim_all.append(sim_i)
  
          # (n_image, n_caption)
          sim_all = torch.cat(sim_all, 1)
  
          return sim_all
  
      def init_weights(self):
          for m in self.children():
              if isinstance(m, nn.Linear):
                  r = np.sqrt(6.) / np.sqrt(m.in_features + m.out_features)
                  m.weight.data.uniform_(-r, r)
                  m.bias.data.fill_(0)
              elif isinstance(m, nn.BatchNorm1d):
                  m.weight.data.fill_(1)
                  m.bias.data.zero_()
  
  
  def SCAN_attention(query, context, smooth, eps=1e-8):
      """
      query: (n_context, queryL, d)
      context: (n_context, sourceL, d)
      """
      # --> (batch, d, queryL)
      queryT = torch.transpose(query, 1, 2)
  
      # (batch, sourceL, d)(batch, d, queryL)
      # --> (batch, sourceL, queryL)
      attn = torch.bmm(context, queryT)
  
      attn = nn.LeakyReLU(0.1)(attn)
      attn = l2norm(attn, 2)
  
      # --> (batch, queryL, sourceL)
      attn = torch.transpose(attn, 1, 2).contiguous()
      # --> (batch, queryL, sourceL
      attn = F.softmax(attn*smooth, dim=2)
  
      # --> (batch, sourceL, queryL)
      attnT = torch.transpose(attn, 1, 2).contiguous()
  
      # --> (batch, d, sourceL)
      contextT = torch.transpose(context, 1, 2)
      # (batch x d x sourceL)(batch x sourceL x queryL)
      # --> (batch, d, queryL)
      weightedContext = torch.bmm(contextT, attnT)
      # --> (batch, queryL, d)
      weightedContext = torch.transpose(weightedContext, 1, 2)
      weightedContext = l2norm(weightedContext, dim=-1)
  
      return weightedContext
  
  
  class SGRAF(nn.Module):
      """
      Similarity Reasoning and Filtration (SGRAF) Network
      """
      def __init__(self,
                   model_name,
                   module_name,
                   sgr_step,
                   embed_size,
                   sim_dim,
                   vocab_size,
                   word_dim,
                   num_layers,
                   image_dim,
                   margin,
                   max_violation,
                   use_bi_gru=True,
                   image_norm=True,
                   text_norm=True,
                   **kwargs):
  
          super(SGRAF, self).__init__()
  
          # Build Models
          self.img_enc = EncoderImage(model_name, image_dim, embed_size, image_norm)
          self.txt_enc = EncoderText(model_name, vocab_size, word_dim, embed_size, num_layers,
                                     use_bi_gru=use_bi_gru, text_norm=text_norm)
          self.sim_enc = EncoderSimilarity(embed_size, sim_dim, module_name, sgr_step)
  
          # Loss and Optimizer
          self.criterion = ContrastiveLoss(margin=margin, max_violation=max_violation)
  
      def forward_emb(self, batch):
          """Compute the image and caption embeddings"""
  
          images = batch['image_feat']
          captions = batch['text_token']
          lengths = batch['text_len']
  
          if torch.cuda.is_available():
              images = images.cuda()
              captions = captions.cuda()
          lengths = lengths.tolist()
  
          # Forward feature encoding
          img_embs = self.img_enc(images)
          cap_embs = self.txt_enc(captions, lengths)
          return img_embs, cap_embs, lengths
  
      def forward_sim(self, batch):
          img_embs, cap_embs, cap_lens = batch
          # Forward similarity encoding
          sims = self.sim_enc(img_embs, cap_embs, cap_lens)
          return sims
  
      def forward(self, batch):
          images = batch['image_feat']
          captions = batch['text_token']
          lengths = batch['text_len']
  
          if torch.cuda.is_available():
              images = images.cuda()
              captions = captions.cuda()
          lengths = lengths.tolist()
  
          img_embs = self.img_enc(images)
          cap_embs = self.txt_enc(captions, lengths)
          sims = self.sim_enc(img_embs, cap_embs, lengths)
          loss = self.criterion(sims)
  
          return loss
  
      @staticmethod
      def cal_sim(model, img_embs, cap_embs, cap_lens, **kwargs):
          shard_size = kwargs.get('shard_size', 100)
  
          n_im_shard = (len(img_embs) - 1) // shard_size + 1
          n_cap_shard = (len(cap_embs) - 1) // shard_size + 1
  
          sims = np.zeros((len(img_embs), len(cap_embs)))
          for i in range(n_im_shard):
              im_start, im_end = shard_size * i, min(shard_size * (i + 1), len(img_embs))
              for j in range(n_cap_shard):
                  ca_start, ca_end = shard_size * j, min(shard_size * (j + 1), len(cap_embs))
  
                  with torch.no_grad():
                      im = torch.FloatTensor(img_embs[im_start:im_end])
                      ca = torch.FloatTensor(cap_embs[ca_start:ca_end])
                      l = cap_lens[ca_start:ca_end].tolist()
  
                      if torch.cuda.is_available():
                          im = im.cuda()
                          ca = ca.cuda()
  
                      sim = model.forward_sim((im, ca, l))
  
                  sims[im_start:im_end, ca_start:ca_end] = sim.cpu().detach().numpy()
          return sims
  
  
apis: torch.nn.functional.softmax torch.sub torch.no_grad torch.cuda.is_available torch.nn.Sigmoid torch.bmm torch.nn.Dropout torch.cat torch.nn.BatchNorm1d torch.nn.Softmax torch.mean torch.transpose torch.FloatTensor torch.nn.Linear torch.nn.Tanh numpy.sqrt torch.nn.ReLU torch.matmul torch.nn.LeakyReLU