class_docstr: Shunted Transformer.
  
  Args:
      img_size (`int`): image size (width or height)
      in_chans (`int`): input channels
      num_classes (`int`): number of classes for classification
      embed_dims (`List[int]`): a list of embedding dimensions at different stages
      num_heads (`List[int]`): a list of attention heads at different stages
      mlp_ratios (`List[int]`): a list of mlp ratios at different stages
      qkv_bias (`bool`): flag to enable bias for qkv
      qk_scale (`int`): number to scale qk
      drop_rate (`float`): drop rate of mlp layers
      attn_drop_rate (`float`): drop rate of attention
      drop_path_rate (`float`): drop rate of drop path
      norm_layer (`nn.Module`): normalization layer
      depths (`List[int]`): a list of stochastic depths for blocks
      sr_ratios (`List[int]`): a list of shunt rates in blocks at each stage
      num_stages (`int`): number of stages
      num_conv (`int`): number of additional convolutions in patch embedding layer at the first stage
class_docstr_tok: ["Shunted","Transformer",".","Args",":","img_size","(","`","int","`",")",":","image","size","(","width","or","height",")","in_chans","(","`","int","`",")",":","input","channels","num_classes","(","`","int","`",")",":","number","of","classes","for","classification","embed_dims","(","`","List","[","int","]","`",")",":","a","list","of","embedding","dimensions","at","different","stages","num_heads","(","`","List","[","int","]","`",")",":","a","list","of","attention","heads","at","different","stages","mlp_ratios","(","`","List","[","int","]","`",")",":","a","list","of","mlp","ratios","at","different","stages","qkv_bias","(","`","bool","`",")",":","flag","to","enable","bias","for","qkv","qk_scale","(","`","int","`",")",":","number","to","scale","qk","drop_rate","(","`","float","`",")",":","drop","rate","of","mlp","layers","attn_drop_rate","(","`","float","`",")",":","drop","rate","of","attention","drop_path_rate","(","`","float","`",")",":","drop","rate","of","drop","path","norm_layer","(","`","nn.Module","`",")",":","normalization","layer","depths","(","`","List","[","int","]","`",")",":","a","list","of","stochastic","depths","for","blocks","sr_ratios","(","`","List","[","int","]","`",")",":","a","list","of","shunt","rates","in","blocks","at","each","stage","num_stages","(","`","int","`",")",":","number","of","stages","num_conv","(","`","int","`",")",":","number","of","additional","convolutions","in","patch","embedding","layer","at","the","first","stage"]
code: def __init__(
      self,
      img_size=224,
      num_classes=1000,
      embed_dims=None,
      num_heads=None,
      mlp_ratios=None,
      qkv_bias=False,
      qk_scale=None,
      drop_rate=0.0,
      attn_drop_rate=0.0,
      drop_path_rate=0.0,
      norm_layer=nn.LayerNorm,
      depths=None,
      sr_ratios=None,
      num_stages=4,
      num_conv=0,
  ):
      super().__init__()
      self.num_classes = num_classes
      self.embed_dims = embed_dims
      self.num_heads = num_heads
      self.mlp_ratios = mlp_ratios
      self.sr_ratios = sr_ratios
      self.depths = depths
      self.num_stages = num_stages
      self.num_conv = num_conv
  
      dpr = [
          x.item()
          for x in torch.linspace(0, drop_path_rate, sum(self.depths))
      ]  # stochastic depth decay rule
      cur = 0
  
      for i in range(self.num_stages):
          if i == 0:
              patch_embed = HeadPatchEmbed(
                  self.num_conv, embed_dim=self.embed_dims[0]
              )
          else:
              patch_embed = OverlapPatchEmbed(
                  img_size=img_size // (2 ** (i + 1)),
                  patch_size=3,
                  stride=2,
                  in_chans=self.embed_dims[i - 1],
                  embed_dim=self.embed_dims[i],
              )
  
          block = nn.ModuleList(
              [
                  Block(
                      dim=self.embed_dims[i],
                      num_heads=self.num_heads[i],
                      mlp_ratio=self.mlp_ratios[i],
                      qkv_bias=qkv_bias,
                      qk_scale=qk_scale,
                      drop=drop_rate,
                      attn_drop=attn_drop_rate,
                      drop_path=dpr[cur + j],
                      norm_layer=norm_layer,
                      sr_ratio=self.sr_ratios[i],
                  )
                  for j in range(self.depths[i])
              ]
          )
          norm = norm_layer(self.embed_dims[i])
          cur += self.depths[i]
  
          setattr(self, f"patch_embed{i + 1}", patch_embed)
          setattr(self, f"block{i + 1}", block)
          setattr(self, f"norm{i + 1}", norm)
  
      # classification head
      self.head = (
          nn.Linear(self.embed_dims[-1], self.num_classes)
          if self.num_classes > 0
          else nn.Identity()
      )
  
      self.apply(init_vit_weights)
  
file_path: towhee/towhee/models/shunted_transformer/shunted_transformer.py
imports: import torch
  from torch import nn
  from towhee.models.utils import create_model
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.shunted_transformer import OverlapPatchEmbed, HeadPatchEmbed, Block, get_configs
label_desc: Use image node_id 1 to create a new ShuntedTransformer object from inherited base classes: nn with example: obj = ShuntedTransformer(img_size, num_classes, embed_dims, num_heads, mlp_ratios, qkv_bias, qk_scale, drop_rate, attn_drop_rate, drop_path_rate, norm_layer, depths, sr_ratios, num_stages, num_conv)