class_docstr: Multiscale Vision Transformers
  Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, Christoph Feichtenhofer
  https://arxiv.org/abs/2104.11227
class_docstr_tok: ["Multiscale","Vision","Transformers","Haoqi","Fan",",","Bo","Xiong",",","Karttikeya","Mangalam",",","Yanghao","Li",",","Zhicheng","Yan",",","Jitendra","Malik",",","Christoph","Feichtenhofer","https",":","\/\/arxiv.org\/abs\/2104.11227"]
code: def __init__(
      self,
      patch_2d: bool,
      patch_stride: List,
      embed_dim: int,
      num_heads: int,
      mlp_ratio: float,
      qkv_bias: bool,
      dropout_rate: float,
      depth: int,
      droppath_rate: float,
      mode: str,
      cls_embed_on: bool,
      sep_pos_embed: bool,
      norm: str,
      patch_kernel: List,
      patch_padding: List,
      pool_q_kernel: List,
      pool_kv_kernel: List,
      pool_skip_kernel: List,
      pool_q_stride: List,
      pool_kv_stride: List,
      pool_skip_stride: List,
      dim_mul_arg: List,
      head_mul_arg: List,
      norm_stem: bool,
      num_classes: int,
      head_act: str,
      train_crop_size: int,
      test_crop_size: int,
      num_frames: int,
      input_channel_num: List,
  ):
      super().__init__()
      assert train_crop_size == test_crop_size
      spatial_size = train_crop_size
      temporal_size = num_frames
      in_chans = input_channel_num[0]
      use_2d_patch = patch_2d
      self.patch_stride = patch_stride
      if use_2d_patch:
          self.patch_stride = [1] + self.patch_stride
      self.drop_rate = dropout_rate
      drop_path_rate = droppath_rate
      self.cls_embed_on = cls_embed_on
      self.sep_pos_embed = sep_pos_embed
      if norm == "layernorm":
          norm_layer = partial(LayerNorm, eps=1e-6)
      else:
          raise NotImplementedError("Only supports layernorm.")
      self.num_classes = num_classes
      self.patch_embed = PatchEmbed(
          dim_in=in_chans,
          dim_out=embed_dim,
          kernel=patch_kernel,
          stride=patch_stride,
          padding=patch_padding,
          conv_2d=use_2d_patch,
      )
      self.input_dims = [temporal_size, spatial_size, spatial_size]
      assert self.input_dims[1] == self.input_dims[2]
      self.patch_dims = [
          self.input_dims[i] // self.patch_stride[i]
          for i in range(len(self.input_dims))
      ]
      num_patches = math.prod(self.patch_dims)
  
      dpr = [
          x.item() for x in torch.linspace(0, drop_path_rate, depth)
      ]  # stochastic depth decay rule
  
      if self.cls_embed_on:
          self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
          pos_embed_dim = num_patches + 1
      else:
          pos_embed_dim = num_patches
  
      if self.sep_pos_embed:
          self.pos_embed_spatial = nn.Parameter(
              torch.zeros(
                  1, self.patch_dims[1] * self.patch_dims[2], embed_dim
              )
          )
          self.pos_embed_temporal = nn.Parameter(
              torch.zeros(1, self.patch_dims[0], embed_dim)
          )
          self.pos_embed_class = nn.Parameter(
              torch.zeros(1, 1, embed_dim)
          )
      else:
          self.pos_embed = nn.Parameter(
              torch.zeros(1, pos_embed_dim, embed_dim)
          )
      self.add = Add()
      if self.drop_rate > 0.0:
          self.pos_drop = nn.Dropout(p=self.drop_rate)
  
      pool_q = pool_q_kernel
      pool_kv = pool_kv_kernel
      pool_skip = pool_skip_kernel
      stride_q = pool_q_stride
      stride_kv = pool_kv_stride
      stride_skip = pool_skip_stride
  
      dim_mul, head_mul = torch.ones(depth + 1), torch.ones(depth + 1)
  
      if len(dim_mul_arg) > 1:
          for k in dim_mul_arg:
              dim_mul[k[0]] = k[1]
      if len(head_mul_arg) > 1:
          for k in head_mul_arg:
              head_mul[k[0]] = k[1]
  
      self.norm_stem = norm_layer(embed_dim) if norm_stem else None
  
      self.blocks = nn.ModuleList()
      for i in range(depth):
          num_heads = round_width(num_heads, head_mul[i])
          embed_dim = round_width(
              embed_dim, dim_mul[i], divisor=num_heads
          )
          dim_out = round_width(
              embed_dim,
              dim_mul[i + 1],
              divisor=round_width(num_heads, head_mul[i + 1]),
          )
          self.blocks.append(
              MultiScaleBlock(
                  dim=embed_dim,
                  dim_out=dim_out,
                  num_heads=num_heads,
                  mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias,
                  drop_rate=self.drop_rate,
                  drop_path=dpr[i],
                  norm_layer=norm_layer,
                  kernel_q=pool_q[i] if len(pool_q) > i else [],
                  kernel_kv=pool_kv[i] if len(pool_kv) > i else [],
                  kernel_skip=pool_skip[i]
                  if len(pool_skip) > i
                  else [],
                  stride_q=stride_q[i] if len(stride_q) > i else [],
                  stride_kv=stride_kv[i] if len(stride_kv) > i else [],
                  stride_skip=stride_skip[i]
                  if len(stride_skip) > i
                  else [],
                  mode=mode,
                  has_cls_embed=self.cls_embed_on,
              )
          )
  
      embed_dim = dim_out
      self.norm = norm_layer(embed_dim)
  
      self.index_select = IndexSelect()
  
      self.head = TransformerBasicHead(
          embed_dim,
          num_classes,
          dropout_rate=dropout_rate,
          act_func=head_act,
      )
      if self.sep_pos_embed:
          trunc_normal_(self.pos_embed_spatial, std=0.02)
          trunc_normal_(self.pos_embed_temporal, std=0.02)
          trunc_normal_(self.pos_embed_class, std=0.02)
      else:
          trunc_normal_(self.pos_embed, std=0.02)
      if self.cls_embed_on:
          trunc_normal_(self.cls_token, std=0.02)
      self.apply(self._init_weights)
      self.out_thw = None
  
      self.num_frames = num_frames
      self.train_crop_size = train_crop_size
  
imports: import math
  from functools import partial
  from typing import Tuple, List
  import numpy
  import torch
  from torch import nn
  from torch.nn.init import trunc_normal_
  from einops import rearrange
  from towhee.models.layers.layers_with_relprop import GELU, Linear, Dropout, Conv2d, Conv3d, Softmax, Sigmoid, LayerNorm, Einsum, Add, Clone, MaxPool3d, IndexSelect
label_desc: Use image node_id 1 to create a new MViT object from inherited base classes: nn with example: obj = MViT(patch_2d, patch_stride, embed_dim, num_heads, mlp_ratio, qkv_bias, dropout_rate, depth, droppath_rate, mode, cls_embed_on, sep_pos_embed, norm, patch_kernel, patch_padding, pool_q_kernel, pool_kv_kernel, pool_skip_kernel, pool_q_stride, pool_kv_stride, pool_skip_stride, dim_mul_arg, head_mul_arg, norm_stem, num_classes, head_act, train_crop_size, test_crop_size, num_frames, input_channel_num)