class_docstr: MobileOne building block.
  
  This block has a multi-branched architecture at train-time
  and plain-CNN style architecture at inference time
  For more details, please refer to our paper:
  `An Improved One millisecond Mobile Backbone` -
  https://arxiv.org/pdf/2206.04040.pdf
class_docstr_tok: ["MobileOne","building","block",".","This","block","has","a","multi-branched","architecture","at","train-time","and","plain-CNN","style","architecture","at","inference","time","For","more","details",",","please","refer","to","our","paper",":","`","An","Improved","One","millisecond","Mobile","Backbone","`","-","https",":","\/\/arxiv.org\/pdf\/2206.04040.pdf"]
docstr: Construct a MobileOneBlock module.
  
  Args:
      in_channels: Number of channels in the input.
      out_channels: Number of channels produced by the block.
      kernel_size: Size of the convolution kernel.
      stride: Stride size.
      padding: Zero-padding size.
      dilation: Kernel dilation factor.
      groups: Group number.
      inference_mode: If True, instantiates model in inference mode.
      use_se: Whether to use SE-ReLU activations.
      use_act: Whether to use activation. Default: ``True``
      use_scale_branch: Whether to use scale branch. Default: ``True``
      num_conv_branches: Number of linear conv branches.
docstr_tok: ["Construct","a","MobileOneBlock","module",".","Args",":","in_channels",":","Number","of","channels","in","the","input",".","out_channels",":","Number","of","channels","produced","by","the","block",".","kernel_size",":","Size","of","the","convolution","kernel",".","stride",":","Stride","size",".","padding",":","Zero-padding","size",".","dilation",":","Kernel","dilation","factor",".","groups",":","Group","number",".","inference_mode",":","If","True",",","instantiates","model","in","inference","mode",".","use_se",":","Whether","to","use","SE-ReLU","activations",".","use_act",":","Whether","to","use","activation",".","Default",":","``","True","``","use_scale_branch",":","Whether","to","use","scale","branch",".","Default",":","``","True","``","num_conv_branches",":","Number","of","linear","conv","branches","."]
code: def __init__(
      self,
      in_channels: int,
      out_channels: int,
      kernel_size: int,
      stride: int = 1,
      padding: int = 0,
      dilation: int = 1,
      groups: int = 1,
      inference_mode: bool = False,
      use_se: bool = False,
      use_act: bool = True,
      use_scale_branch: bool = True,
      num_conv_branches: int = 1,
      activation: nn.Module = nn.GELU(),
  ) -> None:
      """Construct a MobileOneBlock module.
  
      Args:
          in_channels: Number of channels in the input.
          out_channels: Number of channels produced by the block.
          kernel_size: Size of the convolution kernel.
          stride: Stride size.
          padding: Zero-padding size.
          dilation: Kernel dilation factor.
          groups: Group number.
          inference_mode: If True, instantiates model in inference mode.
          use_se: Whether to use SE-ReLU activations.
          use_act: Whether to use activation. Default: ``True``
          use_scale_branch: Whether to use scale branch. Default: ``True``
          num_conv_branches: Number of linear conv branches.
      """
      super(MobileOneBlock, self).__init__()
      self.inference_mode = inference_mode
      self.groups = groups
      self.stride = stride
      self.padding = padding
      self.dilation = dilation
      self.kernel_size = kernel_size
      self.in_channels = in_channels
      self.out_channels = out_channels
      self.num_conv_branches = num_conv_branches
  
      # Check if SE-ReLU is requested
      if use_se:
          self.se = SEBlock(out_channels)
      else:
          self.se = nn.Identity()
  
      if use_act:
          self.activation = activation
      else:
          self.activation = nn.Identity()
  
      if inference_mode:
          self.reparam_conv = nn.Conv2d(
              in_channels=in_channels,
              out_channels=out_channels,
              kernel_size=kernel_size,
              stride=stride,
              padding=padding,
              dilation=dilation,
              groups=groups,
              bias=True,
          )
      else:
          # Re-parameterizable skip connection
          self.rbr_skip = (
              nn.BatchNorm2d(num_features=in_channels)
              if out_channels == in_channels and stride == 1
              else None
          )
  
          # Re-parameterizable conv branches
          if num_conv_branches > 0:
              rbr_conv = list()
              for _ in range(self.num_conv_branches):
                  rbr_conv.append(
                      self._conv_bn(
                          kernel_size=kernel_size, padding=padding
                      )
                  )
              self.rbr_conv = nn.ModuleList(rbr_conv)
          else:
              self.rbr_conv = None
  
          # Re-parameterizable scale branch
          self.rbr_scale = None
          if (kernel_size > 1) and use_scale_branch:
              self.rbr_scale = self._conv_bn(kernel_size=1, padding=0)
  
file_path: ml-fastvit/models/modules/mobileone.py
imports: from typing import Union, Tuple
  import copy
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new MobileOneBlock object from inherited base classes: nn with example: obj = MobileOneBlock(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, inference_mode, use_se, use_act, use_scale_branch, num_conv_branches, activation)