class_docstr: TinyViT model, as described in https://arxiv.org/abs/2207.10666
  
  Args:
      img_size: Size of input image.
      in_chans: Number of input image's channels.
      num_classes: Number of output classes.
      embed_dims: List of embedding dimensions.
      depths: List of block count for each downsampling stage
      num_heads: List of attention heads used in self-attention for each downsampling stage.
      window_sizes: List of self-attention's window size for each downsampling stage.
      mlp_ratio: Ratio of MLP dimension to embedding dimension in self-attention.
      drop_rate: Dropout rate.
      drop_path_rate: Stochastic depth rate.
      use_checkpoint: Whether to use activation checkpointing to trade compute for memory.
      mbconv_expand_ratio: Expansion ratio used in MBConv block.
      local_conv_size: Kernel size of convolution used in TinyViTBlock
      activation: activation function.
      mobile_same: Whether to use modifications for MobileSAM.
class_docstr_tok: ["TinyViT","model",",","as","described","in","https",":","\/\/arxiv.org\/abs\/2207.10666","Args",":","img_size",":","Size","of","input","image",".","in_chans",":","Number","of","input","image","'s","channels",".","num_classes",":","Number","of","output","classes",".","embed_dims",":","List","of","embedding","dimensions",".","depths",":","List","of","block","count","for","each","downsampling","stage","num_heads",":","List","of","attention","heads","used","in","self-attention","for","each","downsampling","stage",".","window_sizes",":","List","of","self-attention","'s","window","size","for","each","downsampling","stage",".","mlp_ratio",":","Ratio","of","MLP","dimension","to","embedding","dimension","in","self-attention",".","drop_rate",":","Dropout","rate",".","drop_path_rate",":","Stochastic","depth","rate",".","use_checkpoint",":","Whether","to","use","activation","checkpointing","to","trade","compute","for","memory",".","mbconv_expand_ratio",":","Expansion","ratio","used","in","MBConv","block",".","local_conv_size",":","Kernel","size","of","convolution","used","in","TinyViTBlock","activation",":","activation","function",".","mobile_same",":","Whether","to","use","modifications","for","MobileSAM","."]
code: def __init__(
      self,
      img_size: int = 224,
      in_chans: int = 3,
      num_classes: int = 1000,
      embed_dims: list[int] = [96, 192, 384, 768],
      depths: list[int] = [2, 2, 6, 2],
      num_heads: list[int] = [3, 6, 12, 24],
      window_sizes: list[int] = [7, 7, 14, 7],
      mlp_ratio: float = 4.0,
      drop_rate: float = 0.0,
      drop_path_rate: float = 0.0,
      use_checkpoint: bool = False,
      mbconv_expand_ratio: float = 4.0,
      local_conv_size: int = 3,
      # layer_lr_decay: float = 1.0,
      activation: type[Module] = nn.GELU,
      mobile_sam: bool = False,
  ) -> None:
      super().__init__()
      self.img_size = img_size
      self.mobile_sam = mobile_sam
      self.neck: Optional[Module]
      if mobile_sam:
          # MobileSAM adjusts the stride to match the total stride of other ViT backbones
          # used in the original SAM (stride 16)
          strides = [2, 2, 1, 1]
          self.neck = nn.Sequential(
              nn.Conv2d(embed_dims[-1], 256, 1, bias=False),
              LayerNorm2d(256),
              nn.Conv2d(256, 256, 3, 1, 1, bias=False),
              LayerNorm2d(256),
          )
      else:
          strides = [2, 2, 2, 1]
          self.neck = None
  
      self.patch_embed = PatchEmbed(in_chans, embed_dims[0], activation)
      input_resolution = img_size // 4
  
      # NOTE: if we don't support training, this might be unimportant
      # stochastic depth decay rule
      dpr = [
          x.item()
          for x in torch.linspace(0, drop_path_rate, sum(depths))
      ]
  
      # build layers
      n_layers = len(depths)
      layers = []
      for i_layer, (
          embed_dim,
          depth,
          num_heads_i,
          window_size,
          stride,
      ) in enumerate(
          zip(embed_dims, depths, num_heads, window_sizes, strides)
      ):
          out_dim = embed_dims[min(i_layer + 1, len(embed_dims) - 1)]
          downsample = (
              PatchMerging(
                  input_resolution,
                  embed_dim,
                  out_dim,
                  stride,
                  activation,
              )
              if (i_layer < n_layers - 1)
              else None
          )
          kwargs: dict[str, Any] = {
              "dim": embed_dim,
              "depth": depth,
              "drop_path": dpr[
                  sum(depths[:i_layer]) : sum(depths[: i_layer + 1])
              ],
              "downsample": downsample,
              "use_checkpoint": use_checkpoint,
              "activation": activation,
          }
          layer: ConvLayer | BasicLayer
          if i_layer == 0:
              layer = ConvLayer(
                  conv_expand_ratio=mbconv_expand_ratio, **kwargs
              )
          else:
              layer = BasicLayer(
                  input_resolution=input_resolution,
                  num_heads=num_heads_i,
                  window_size=window_size,
                  mlp_ratio=mlp_ratio,
                  drop=drop_rate,
                  local_conv_size=local_conv_size,
                  **kwargs,
              )
          layers.append(layer)
          input_resolution //= stride
      self.layers = nn.Sequential(*layers)
      self.feat_size = input_resolution  # final feature map size
  
      # Classifier head
      # NOTE: this is redundant for MobileSAM, but we still need it
      # to load pre-trained weights with strict=True
      # TODO: enable strict=False, or host our own weights
      self.norm_head = nn.LayerNorm(embed_dims[-1])
      self.head = nn.Linear(embed_dims[-1], num_classes)
  
imports: from __future__ import annotations
  import itertools
  import warnings
  from typing import Any, Optional
  import torch
  import torch.nn.functional
  from torch import nn
  from torch.utils import checkpoint
  from kornia.contrib.models.common import DropPath, LayerNorm2d, window_partition, window_unpartition
  from kornia.core import Module, Tensor
  from kornia.core.check import KORNIA_CHECK
label_desc: Use image node_id 1 to create a new TinyViT object from inherited base classes: Module with example: obj = TinyViT(img_size, in_chans, num_classes, embed_dims, depths, num_heads, window_sizes, mlp_ratio, drop_rate, drop_path_rate, use_checkpoint, mbconv_expand_ratio, local_conv_size, activation, mobile_sam)