class_docstr: This class implements `FastViT architecture <https://arxiv.org/pdf/2303.14189.pdf>`_
class_docstr_tok: ["This","class","implements","`","FastViT","architecture","<","https",":","\/\/arxiv.org\/pdf\/2303.14189.pdf",">","`","_"]
code: def __init__(
      self,
      layers,
      token_mixers: Tuple[str, ...],
      embed_dims=None,
      mlp_ratios=None,
      downsamples=None,
      repmixer_kernel_size=3,
      norm_layer: nn.Module = nn.BatchNorm2d,
      act_layer: nn.Module = nn.GELU,
      num_classes=1000,
      pos_embs=None,
      down_patch_size=7,
      down_stride=2,
      drop_rate=0.0,
      drop_path_rate=0.0,
      use_layer_scale=True,
      layer_scale_init_value=1e-5,
      fork_feat=False,
      init_cfg=None,
      pretrained=None,
      cls_ratio=2.0,
      inference_mode=False,
      **kwargs,
  ) -> None:
      super().__init__()
  
      if not fork_feat:
          self.num_classes = num_classes
      self.fork_feat = fork_feat
  
      if pos_embs is None:
          pos_embs = [None] * len(layers)
  
      # Convolutional stem
      self.patch_embed = convolutional_stem(
          3, embed_dims[0], inference_mode
      )
  
      # Build the main stages of the network architecture
      network = []
      for i in range(len(layers)):
          # Add position embeddings if requested
          if pos_embs[i] is not None:
              network.append(
                  pos_embs[i](
                      embed_dims[i],
                      embed_dims[i],
                      inference_mode=inference_mode,
                  )
              )
          stage = basic_blocks(
              embed_dims[i],
              i,
              layers,
              token_mixer_type=token_mixers[i],
              kernel_size=repmixer_kernel_size,
              mlp_ratio=mlp_ratios[i],
              act_layer=act_layer,
              norm_layer=norm_layer,
              drop_rate=drop_rate,
              drop_path_rate=drop_path_rate,
              use_layer_scale=use_layer_scale,
              layer_scale_init_value=layer_scale_init_value,
              inference_mode=inference_mode,
          )
          network.append(stage)
          if i >= len(layers) - 1:
              break
  
          # Patch merging/downsampling between stages.
          if downsamples[i] or embed_dims[i] != embed_dims[i + 1]:
              network.append(
                  PatchEmbed(
                      patch_size=down_patch_size,
                      stride=down_stride,
                      in_channels=embed_dims[i],
                      embed_dim=embed_dims[i + 1],
                      inference_mode=inference_mode,
                  )
              )
  
      self.network = nn.ModuleList(network)
  
      # For segmentation and detection, extract intermediate output
      if self.fork_feat:
          # add a norm layer for each output
          self.out_indices = [0, 2, 4, 6]
          for i_emb, i_layer in enumerate(self.out_indices):
              if i_emb == 0 and os.environ.get("FORK_LAST3", None):
                  """For RetinaNet, `start_level=1`. The first norm layer will not used.
                  cmd: `FORK_LAST3=1 python -m torch.distributed.launch ...`
                  """
                  layer = nn.Identity()
              else:
                  layer = norm_layer(embed_dims[i_emb])
              layer_name = f"norm{i_layer}"
              self.add_module(layer_name, layer)
      else:
          # Classifier head
          self.gap = nn.AdaptiveAvgPool2d(output_size=1)
          self.conv_exp = MobileOneBlock(
              in_channels=embed_dims[-1],
              out_channels=int(embed_dims[-1] * cls_ratio),
              kernel_size=3,
              stride=1,
              padding=1,
              groups=embed_dims[-1],
              inference_mode=inference_mode,
              use_se=True,
              num_conv_branches=1,
          )
          self.head = (
              nn.Linear(int(embed_dims[-1] * cls_ratio), num_classes)
              if num_classes > 0
              else nn.Identity()
          )
  
      self.apply(self.cls_init_weights)
      self.init_cfg = copy.deepcopy(init_cfg)
  
      # load pre-trained model
      if self.fork_feat and (
          self.init_cfg is not None or pretrained is not None
      ):
          self.init_weights()
  
imports: import os
  import copy
  from functools import partial
  from typing import List, Tuple, Optional, Union
  import torch
  import torch.nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.models.layers import DropPath, trunc_normal_
  from timm.models.registry import register_model
  from models.modules.mobileone import MobileOneBlock
  from models.modules.replknet import ReparamLargeKernelConv
label_desc: Use image node_id 1 to create a new FastViT object from inherited base classes: nn with example: obj = FastViT(layers, token_mixers, embed_dims, mlp_ratios, downsamples, repmixer_kernel_size, norm_layer, act_layer, num_classes, pos_embs, down_patch_size, down_stride, drop_rate, drop_path_rate, use_layer_scale, layer_scale_init_value, fork_feat, init_cfg, pretrained, cls_ratio, inference_mode)