class_docstr: Improved Multiscale Vision Transformers for Classification and Detection
  Yanghao Li*, Chao-Yuan Wu*, Haoqi Fan, Karttikeya Mangalam, Bo Xiong, Jitendra Malik,
      Christoph Feichtenhofer*
  https://arxiv.org/abs/2112.01526
  
  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: ["Improved","Multiscale","Vision","Transformers","for","Classification","and","Detection","Yanghao","Li","*",",","Chao-Yuan","Wu","*",",","Haoqi","Fan",",","Karttikeya","Mangalam",",","Bo","Xiong",",","Jitendra","Malik",",","Christoph","Feichtenhofer","*","https",":","\/\/arxiv.org\/abs\/2112.01526","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,
      cfg: MultiScaleVitCfg,
      img_size: Tuple[int, int] = (224, 224),
      in_chans: int = 3,
      global_pool: Optional[str] = None,
      num_classes: int = 1000,
      drop_path_rate: float = 0.0,
      drop_rate: float = 0.0,
  ):
      super().__init__()
      img_size = to_2tuple(img_size)
      norm_layer = partial(
          get_norm_layer(cfg.norm_layer), eps=cfg.norm_eps
      )
      self.num_classes = num_classes
      self.drop_rate = drop_rate
      if global_pool is None:
          global_pool = "token" if cfg.use_cls_token else "avg"
      self.global_pool = global_pool
      self.depths = tuple(cfg.depths)
      self.expand_attn = cfg.expand_attn
  
      embed_dim = cfg.embed_dim[0]
      self.patch_embed = PatchEmbed(
          dim_in=in_chans,
          dim_out=embed_dim,
          kernel=cfg.patch_kernel,
          stride=cfg.patch_stride,
          padding=cfg.patch_padding,
      )
      patch_dims = (
          img_size[0] // cfg.patch_stride[0],
          img_size[1] // cfg.patch_stride[1],
      )
      num_patches = prod(patch_dims)
  
      if cfg.use_cls_token:
          self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
          self.num_prefix_tokens = 1
          pos_embed_dim = num_patches + 1
      else:
          self.num_prefix_tokens = 0
          self.cls_token = None
          pos_embed_dim = num_patches
  
      if cfg.use_abs_pos:
          self.pos_embed = nn.Parameter(
              torch.zeros(1, pos_embed_dim, embed_dim)
          )
      else:
          self.pos_embed = None
  
      num_stages = len(cfg.embed_dim)
      feat_size = patch_dims
      dpr = [
          x.tolist()
          for x in torch.linspace(
              0, drop_path_rate, sum(cfg.depths)
          ).split(cfg.depths)
      ]
      self.stages = nn.ModuleList()
      for i in range(num_stages):
          if cfg.expand_attn:
              dim_out = cfg.embed_dim[i]
          else:
              dim_out = cfg.embed_dim[min(i + 1, num_stages - 1)]
          stage = MultiScaleVitStage(
              dim=embed_dim,
              dim_out=dim_out,
              depth=cfg.depths[i],
              num_heads=cfg.num_heads[i],
              feat_size=feat_size,
              mlp_ratio=cfg.mlp_ratio,
              qkv_bias=cfg.qkv_bias,
              mode=cfg.mode,
              pool_first=cfg.pool_first,
              expand_attn=cfg.expand_attn,
              kernel_q=cfg.kernel_qkv,
              kernel_kv=cfg.kernel_qkv,
              stride_q=cfg.stride_q[i],
              stride_kv=cfg.stride_kv[i],
              has_cls_token=cfg.use_cls_token,
              rel_pos_type=cfg.rel_pos_type,
              residual_pooling=cfg.residual_pooling,
              norm_layer=norm_layer,
              drop_path=dpr[i],
          )
          embed_dim = dim_out
          feat_size = stage.feat_size
          self.stages.append(stage)
  
      self.num_features = embed_dim
      self.norm = norm_layer(embed_dim)
      self.head = nn.Sequential(
          OrderedDict(
              [
                  ("drop", nn.Dropout(self.drop_rate)),
                  (
                      "fc",
                      nn.Linear(self.num_features, num_classes)
                      if num_classes > 0
                      else nn.Identity(),
                  ),
              ]
          )
      )
  
      if self.pos_embed is not None:
          trunc_normal_tf_(self.pos_embed, std=0.02)
      if self.cls_token is not None:
          trunc_normal_tf_(self.cls_token, std=0.02)
      self.apply(self._init_weights)
  
imports: import operator
  from collections import OrderedDict
  from dataclasses import dataclass
  from functools import partial, reduce
  from typing import Union, List, Tuple, Optional
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import Mlp, DropPath, trunc_normal_tf_, get_norm_layer, to_2tuple
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._registry import register_model, register_model_deprecations, generate_default_cfgs
label_desc: Use image node_id 1 to create a new MultiScaleVit object from inherited base classes: nn with example: obj = MultiScaleVit(cfg, img_size, in_chans, global_pool, num_classes, drop_path_rate, drop_rate)