class_docstr: Expectation Maximization Attention Networks for Semantic Segmentation.
  
  This head is the implementation of `EMANet
  <https://arxiv.org/abs/1907.13426>`_.
  
  Args:
      ema_channels (int): EMA module channels
      num_bases (int): Number of bases.
      num_stages (int): Number of the EM iterations.
      concat_input (bool): Whether concat the input and output of convs
          before classification layer. Default: True
      momentum (float): Momentum to update the base. Default: 0.1.
class_docstr_tok: ["Expectation","Maximization","Attention","Networks","for","Semantic","Segmentation",".","This","head","is","the","implementation","of","`","EMANet","<","https",":","\/\/arxiv.org\/abs\/1907.13426",">","`","_",".","Args",":","ema_channels","(","int",")",":","EMA","module","channels","num_bases","(","int",")",":","Number","of","bases",".","num_stages","(","int",")",":","Number","of","the","EM","iterations",".","concat_input","(","bool",")",":","Whether","concat","the","input","and","output","of","convs","before","classification","layer",".","Default",":","True","momentum","(","float",")",":","Momentum","to","update","the","base",".","Default",":","0.1","."]
code: def __init__(
      self,
      ema_channels,
      num_bases,
      num_stages,
      concat_input=True,
      momentum=0.1,
      **kwargs,
  ):
      super(EMAHead, self).__init__(**kwargs)
      self.ema_channels = ema_channels
      self.num_bases = num_bases
      self.num_stages = num_stages
      self.concat_input = concat_input
      self.momentum = momentum
      self.ema_module = EMAModule(
          self.ema_channels,
          self.num_bases,
          self.num_stages,
          self.momentum,
      )
  
      self.ema_in_conv = ConvModule(
          self.in_channels,
          self.ema_channels,
          3,
          padding=1,
          conv_cfg=self.conv_cfg,
          norm_cfg=self.norm_cfg,
          act_cfg=self.act_cfg,
      )
      # project (0, inf) -> (-inf, inf)
      self.ema_mid_conv = ConvModule(
          self.ema_channels,
          self.ema_channels,
          1,
          conv_cfg=self.conv_cfg,
          norm_cfg=None,
          act_cfg=None,
      )
      for param in self.ema_mid_conv.parameters():
          param.requires_grad = False
  
      self.ema_out_conv = ConvModule(
          self.ema_channels,
          self.ema_channels,
          1,
          conv_cfg=self.conv_cfg,
          norm_cfg=self.norm_cfg,
          act_cfg=None,
      )
      self.bottleneck = ConvModule(
          self.ema_channels,
          self.channels,
          3,
          padding=1,
          conv_cfg=self.conv_cfg,
          norm_cfg=self.norm_cfg,
          act_cfg=self.act_cfg,
      )
      if self.concat_input:
          self.conv_cat = ConvModule(
              self.in_channels + self.channels,
              self.channels,
              kernel_size=3,
              padding=1,
              conv_cfg=self.conv_cfg,
              norm_cfg=self.norm_cfg,
              act_cfg=self.act_cfg,
          )
  
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 EMAHead object from inherited base classes: BaseDecodeHead with example: obj = EMAHead(ema_channels, num_bases, num_stages, concat_input, momentum)