class_docstr: Expectation Maximization Attention Module used in EMANet.
  
  Args:
      channels (int): Channels of the whole module.
      num_bases (int): Number of bases.
      num_stages (int): Number of the EM iterations.
class_docstr_tok: ["Expectation","Maximization","Attention","Module","used","in","EMANet",".","Args",":","channels","(","int",")",":","Channels","of","the","whole","module",".","num_bases","(","int",")",":","Number","of","bases",".","num_stages","(","int",")",":","Number","of","the","EM","iterations","."]
docstr: Forward function.
docstr_tok: ["Forward","function","."]
code: def forward(self, feats):
      """Forward function."""
      batch_size, channels, height, width = feats.size()
      # [batch_size, channels, height*width]
      feats = feats.view(batch_size, channels, height * width)
      # [batch_size, channels, num_bases]
      bases = self.bases.repeat(batch_size, 1, 1)
  
      with torch.no_grad():
          for i in range(self.num_stages):
              # [batch_size, height*width, num_bases]
              attention = torch.einsum("bcn,bck->bnk", feats, bases)
              attention = F.softmax(attention, dim=2)
              # l1 norm
              attention_normed = F.normalize(attention, dim=1, p=1)
              # [batch_size, channels, num_bases]
              bases = torch.einsum(
                  "bcn,bnk->bck", feats, attention_normed
              )
              # l2 norm
              bases = F.normalize(bases, dim=1, p=2)
  
      feats_recon = torch.einsum("bck,bnk->bcn", bases, attention)
      feats_recon = feats_recon.view(
          batch_size, channels, height, width
      )
  
      if self.training:
          bases = bases.mean(dim=0, keepdim=True)
          bases = reduce_mean(bases)
          # l2 norm
          bases = F.normalize(bases, dim=1, p=2)
          self.bases = (
              1 - self.momentum
          ) * self.bases + self.momentum * bases
  
      return feats_recon
  
file_path: lavis/lavis/common/annotator/uniformer/mmseg/models/decode_heads/ema_head.py
imports: import math
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from annotator.uniformer.mmcv.cnn import ConvModule
  from ..builder import HEADS
  from .decode_head import BaseDecodeHead
label_desc: Use image node_id 2 for calling the EMAModule obj's underlying member method code with example usage: obj.forward(feats) and returns: feats_recon