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","."]
code: def __init__(self, channels, num_bases, num_stages, momentum):
      super(EMAModule, self).__init__()
      assert num_stages >= 1, "num_stages must be at least 1!"
      self.num_bases = num_bases
      self.num_stages = num_stages
      self.momentum = momentum
  
      bases = torch.zeros(1, channels, self.num_bases)
      bases.normal_(0, math.sqrt(2.0 / self.num_bases))
      # [1, channels, num_bases]
      bases = F.normalize(bases, dim=1, p=2)
      self.register_buffer("bases", bases)
  
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 1 to create a new EMAModule object from inherited base classes: nn with example: obj = EMAModule(channels, num_bases, num_stages, momentum)