class_docstr: WaveViT
  Args:
      in_chans (`int`):
          Number of image channels.
      num_classes (`int`):
          Number of image classes.
      stem_hidden_dim (`int`):
          Stem layer hidden dim.
      embed_dims (`list`):
          Every layer`s embedding dim.
      num_heads (`list`):
          Every layer`s head number.
      mlp_ratios (`list`):
          Every layer`s mlp ration.
      drop_path_rate (`float`):
          Every layer`s drop rate.
      norm_layer (`nn.Module`):
          Default nn.LayerNorm.
      depths (`int`):
          Number of depths.
      sr_ratios (`list`):
          Every block`s sr ration.
      num_stages (`int`):
          Number of stages.
      token_label (`bool`):
          Whether to use the token_label.
class_docstr_tok: ["WaveViT","Args",":","in_chans","(","`","int","`",")",":","Number","of","image","channels",".","num_classes","(","`","int","`",")",":","Number","of","image","classes",".","stem_hidden_dim","(","`","int","`",")",":","Stem","layer","hidden","dim",".","embed_dims","(","`","list","`",")",":","Every","layer","`","s","embedding","dim",".","num_heads","(","`","list","`",")",":","Every","layer","`","s","head","number",".","mlp_ratios","(","`","list","`",")",":","Every","layer","`","s","mlp","ration",".","drop_path_rate","(","`","float","`",")",":","Every","layer","`","s","drop","rate",".","norm_layer","(","`","nn.Module","`",")",":","Default","nn.LayerNorm",".","depths","(","`","int","`",")",":","Number","of","depths",".","sr_ratios","(","`","list","`",")",":","Every","block","`","s","sr","ration",".","num_stages","(","`","int","`",")",":","Number","of","stages",".","token_label","(","`","bool","`",")",":","Whether","to","use","the","token_label","."]
code: def __init__(
      self,
      in_chans: int = 3,
      num_classes: int = 1000,
      stem_hidden_dim: int = 32,
      embed_dims: List = None,
      num_heads: List = None,
      mlp_ratios: List = None,
      drop_path_rate: float = 0.0,
      norm_layer: Any = nn.LayerNorm,
      depths: int = None,
      sr_ratios: List = None,
      num_stages: int = 4,
      token_label: bool = True,
  ):
      super().__init__()
      if embed_dims is None:
          embed_dims = [64, 128, 320, 448]
      if num_heads is None:
          num_heads = [2, 4, 10, 14]
      if mlp_ratios is None:
          mlp_ratios = [8, 8, 4, 4]
      if depths is None:
          depths = [3, 4, 6, 3]
      if sr_ratios is None:
          sr_ratios = [4, 2, 1, 1]
      self.num_classes = num_classes
      self.depths = depths
      self.num_stages = num_stages
  
      dpr = [
          x.item()
          for x in torch.linspace(0, drop_path_rate, sum(depths))
      ]  # stochastic depth decay rule
      cur = 0
  
      for i in range(num_stages):
          if i == 0:
              patch_embed = Stem(
                  in_chans, stem_hidden_dim, embed_dims[i]
              )
          else:
              patch_embed = DownSamples(
                  embed_dims[i - 1], embed_dims[i]
              )
  
          block = nn.ModuleList(
              [
                  Block(
                      dim=embed_dims[i],
                      num_heads=num_heads[i],
                      mlp_ratio=mlp_ratios[i],
                      drop_path=dpr[cur + j],
                      norm_layer=norm_layer,
                      sr_ratio=sr_ratios[i],
                      block_type="wave" if i < 2 else "std_att",
                  )
                  for j in range(depths[i])
              ]
          )
  
          norm = norm_layer(embed_dims[i])
          cur += 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)
  
      post_layers = ["ca"]
      self.post_network = nn.ModuleList(
          [
              ClassBlock(
                  dim=embed_dims[-1],
                  num_heads=num_heads[-1],
                  mlp_ratio=mlp_ratios[-1],
                  norm_layer=norm_layer,
              )
              for _ in range(len(post_layers))
          ]
      )
  
      # classification head
      self.head = (
          nn.Linear(embed_dims[-1], num_classes)
          if num_classes > 0
          else nn.Identity()
      )
      # token_label
      self.return_dense = token_label
      self.mix_token = token_label
      self.beta = 1.0
      self.pooling_scale = 8
      if self.return_dense:
          self.aux_head = (
              nn.Linear(embed_dims[-1], num_classes)
              if num_classes > 0
              else nn.Identity()
          )
      # token_label
  
      self.apply(init_vit_weights)
  
file_path: towhee/towhee/models/wave_vit/wave_vit.py
imports: import torch
  from torch import nn
  from typing import List, Any
  import numpy
  from timm.models.vision_transformer import _cfg
  from towhee.models.wave_vit.wave_vit_block import rand_bbox, Stem, DownSamples, Block, ClassBlock
  from towhee.models.utils.init_vit_weights import init_vit_weights
  from towhee.models.wave_vit.wave_vit_utils import get_configs
label_desc: Use image node_id 1 to create a new WaveViT object from inherited base classes: nn with example: obj = WaveViT(in_chans, num_classes, stem_hidden_dim, embed_dims, num_heads, mlp_ratios, drop_path_rate, norm_layer, depths, sr_ratios, num_stages, token_label)