class_docstr: This class implements the `MobileViT architecture <https://arxiv.org/abs/2110.02178?context=cs.LG>`_
class_docstr_tok: ["This","class","implements","the","`","MobileViT","architecture","<","https",":","\/\/arxiv.org\/abs\/2110.02178","?","context=cs.LG",">","`","_"]
code: def __init__(self, model_cfg: Dict, num_classes: int = 1000):
      super().__init__()
  
      image_channels = 3
      out_channels = 16
  
      self.conv_1 = ConvLayer(
          in_channels=image_channels,
          out_channels=out_channels,
          kernel_size=3,
          stride=2,
      )
  
      self.layer_1, out_channels = self._make_layer(
          input_channel=out_channels, cfg=model_cfg["layer1"]
      )
      self.layer_2, out_channels = self._make_layer(
          input_channel=out_channels, cfg=model_cfg["layer2"]
      )
      self.layer_3, out_channels = self._make_layer(
          input_channel=out_channels, cfg=model_cfg["layer3"]
      )
      self.layer_4, out_channels = self._make_layer(
          input_channel=out_channels, cfg=model_cfg["layer4"]
      )
      self.layer_5, out_channels = self._make_layer(
          input_channel=out_channels, cfg=model_cfg["layer5"]
      )
  
      exp_channels = min(
          model_cfg["last_layer_exp_factor"] * out_channels, 960
      )
      self.conv_1x1_exp = ConvLayer(
          in_channels=out_channels,
          out_channels=exp_channels,
          kernel_size=1,
      )
  
      self.classifier = nn.Sequential()
      self.classifier.add_module(
          name="global_pool", module=nn.AdaptiveAvgPool2d(1)
      )
      self.classifier.add_module(name="flatten", module=nn.Flatten())
      if 0.0 < model_cfg["cls_dropout"] < 1.0:
          self.classifier.add_module(
              name="dropout",
              module=nn.Dropout(p=model_cfg["cls_dropout"]),
          )
      self.classifier.add_module(
          name="fc",
          module=nn.Linear(
              in_features=exp_channels, out_features=num_classes
          ),
      )
  
      # weight init
      self.apply(self.init_parameters)
  
imports: from typing import Optional, Tuple, Union, Dict
  import math
  import torch
  import torch.nn
  from torch import Tensor
  from torch.nn import functional
  from transformer import TransformerEncoder
  from model_config import get_config
label_desc: Use image node_id 1 to create a new MobileViT object from inherited base classes: nn with example: obj = MobileViT(model_cfg, num_classes)