class_docstr: Module MobileViT. Default arguments is for MobileViT XXS.
  
  Paper: https://arxiv.org/abs/2110.02178
  Based on: https://github.com/chinhsuanwu/mobilevit-pytorch
  
  Args:
      mode: 'xxs', 'xs' or 's', defaults to 'xxs'.
      in_channels: the number of channels for the input image.
      patch_size: image_size must be divisible by patch_size.
      dropout: dropout ratio in Transformer.
  
  Example:
      >>> img = torch.rand(1, 3, 256, 256)
      >>> mvit = MobileViT(mode='xxs')
      >>> mvit(img).shape
      torch.Size([1, 320, 8, 8])
class_docstr_tok: ["Module","MobileViT",".","Default","arguments","is","for","MobileViT","XXS",".","Paper",":","https",":","\/\/arxiv.org\/abs\/2110.02178","Based","on",":","https",":","\/\/github.com\/chinhsuanwu\/mobilevit-pytorch","Args",":","mode",":","'xxs","'",",","'xs","'","or","'s","'",",","defaults","to","'xxs","'",".","in_channels",":","the","number","of","channels","for","the","input","image",".","patch_size",":","image_size","must","be","divisible","by","patch_size",".","dropout",":","dropout","ratio","in","Transformer",".","Example",":",">",">",">","img","=","torch.rand","(","1",",","3",",","256",",","256",")",">",">",">","mvit","=","MobileViT","(","mode='xxs","'",")",">",">",">","mvit","(","img",")",".shape","torch.Size","(","[","1",",","320",",","8",",","8","]",")"]
code: def __init__(
      self,
      mode: str = "xxs",
      in_channels: int = 3,
      patch_size: Tuple[int, int] = (2, 2),
      dropout: float = 0.0,
  ) -> None:
      super().__init__()
      if mode == "xxs":
          expansion = 2
          dims = [64, 80, 96]
          channels = [16, 16, 24, 24, 48, 48, 64, 64, 80, 80, 320]
      elif mode == "xs":
          expansion = 4
          dims = [96, 120, 144]
          channels = [16, 32, 48, 48, 64, 64, 80, 80, 96, 96, 384]
      elif mode == "s":
          expansion = 4
          dims = [144, 192, 240]
          channels = [16, 32, 64, 64, 96, 96, 128, 128, 160, 160, 640]
  
      kernel_size = 3
      depth = [2, 4, 3]
  
      self.conv1 = conv_nxn_bn(in_channels, channels[0], stride=2)
  
      self.mv2 = nn.ModuleList([])
      self.mv2.append(MV2Block(channels[0], channels[1], 1, expansion))
      self.mv2.append(MV2Block(channels[1], channels[2], 2, expansion))
      self.mv2.append(MV2Block(channels[2], channels[3], 1, expansion))
      self.mv2.append(
          MV2Block(channels[2], channels[3], 1, expansion)
      )  # Repeat
      self.mv2.append(MV2Block(channels[3], channels[4], 2, expansion))
      self.mv2.append(MV2Block(channels[5], channels[6], 2, expansion))
      self.mv2.append(MV2Block(channels[7], channels[8], 2, expansion))
  
      self.mvit = nn.ModuleList([])
      self.mvit.append(
          MobileViTBlock(
              dims[0],
              depth[0],
              channels[5],
              kernel_size,
              patch_size,
              int(dims[0] * 2),
              dropout=dropout,
          )
      )
      self.mvit.append(
          MobileViTBlock(
              dims[1],
              depth[1],
              channels[7],
              kernel_size,
              patch_size,
              int(dims[1] * 4),
              dropout=dropout,
          )
      )
      self.mvit.append(
          MobileViTBlock(
              dims[2],
              depth[2],
              channels[9],
              kernel_size,
              patch_size,
              int(dims[2] * 4),
              dropout=dropout,
          )
      )
  
      self.conv2 = conv_1x1_bn(channels[-2], channels[-1])
  
imports: from typing import Any, Dict, Tuple
  import torch
  from torch import nn
  from kornia.core import Module, Tensor
label_desc: Use image node_id 1 to create a new MobileViT object from inherited base classes: Module with example: obj = MobileViT(mode, in_channels, patch_size, dropout)