class_docstr: DaViT
      A PyTorch implementation of `DaViT: Dual Attention Vision Transformers`  - https://arxiv.org/abs/2204.03645
      Supports arbitrary input sizes and pyramid feature extraction
      
  Args:
      in_chans (int): Number of input image channels. Default: 3
      num_classes (int): Number of classes for classification head. Default: 1000
      depths (tuple(int)): Number of blocks in each stage. Default: (1, 1, 3, 1)
      embed_dims (tuple(int)): Patch embedding dimension. Default: (96, 192, 384, 768)
      num_heads (tuple(int)): Number of attention heads in different layers. Default: (3, 6, 12, 24)
      window_size (int): Window size. Default: 7
      mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
      qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
      drop_path_rate (float): Stochastic depth rate. Default: 0.1
      norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
class_docstr_tok: ["DaViT","A","PyTorch","implementation","of","`","DaViT",":","Dual","Attention","Vision","Transformers","`","-","https",":","\/\/arxiv.org\/abs\/2204.03645","Supports","arbitrary","input","sizes","and","pyramid","feature","extraction","Args",":","in_chans","(","int",")",":","Number","of","input","image","channels",".","Default",":","3","num_classes","(","int",")",":","Number","of","classes","for","classification","head",".","Default",":","1000","depths","(","tuple","(","int",")",")",":","Number","of","blocks","in","each","stage",".","Default",":","(","1",",","1",",","3",",","1",")","embed_dims","(","tuple","(","int",")",")",":","Patch","embedding","dimension",".","Default",":","(","96",",","192",",","384",",","768",")","num_heads","(","tuple","(","int",")",")",":","Number","of","attention","heads","in","different","layers",".","Default",":","(","3",",","6",",","12",",","24",")","window_size","(","int",")",":","Window","size",".","Default",":","7","mlp_ratio","(","float",")",":","Ratio","of","mlp","hidden","dim","to","embedding","dim",".","Default",":","4","qkv_bias","(","bool",")",":","If","True",",","add","a","learnable","bias","to","query",",","key",",","value",".","Default",":","True","drop_path_rate","(","float",")",":","Stochastic","depth","rate",".","Default",":","0.1","norm_layer","(","nn.Module",")",":","Normalization","layer",".","Default",":","nn.LayerNorm","."]
code: def __init__(
      self,
      in_chans=3,
      depths=(1, 1, 3, 1),
      embed_dims=(96, 192, 384, 768),
      num_heads=(3, 6, 12, 24),
      window_size=7,
      mlp_ratio=4,
      qkv_bias=True,
      norm_layer="layernorm2d",
      norm_layer_cl="layernorm",
      norm_eps=1e-5,
      attn_types=("spatial", "channel"),
      ffn=True,
      cpe_act=False,
      drop_rate=0.0,
      drop_path_rate=0.0,
      num_classes=1000,
      global_pool="avg",
      head_norm_first=False,
  ):
      super().__init__()
      num_stages = len(embed_dims)
      assert num_stages == len(num_heads) == len(depths)
      norm_layer = partial(get_norm_layer(norm_layer), eps=norm_eps)
      norm_layer_cl = partial(
          get_norm_layer(norm_layer_cl), eps=norm_eps
      )
      self.num_classes = num_classes
      self.num_features = embed_dims[-1]
      self.drop_rate = drop_rate
      self.grad_checkpointing = False
      self.feature_info = []
  
      self.stem = Stem(in_chans, embed_dims[0], norm_layer=norm_layer)
      in_chs = embed_dims[0]
  
      dpr = [
          x.tolist()
          for x in torch.linspace(0, drop_path_rate, sum(depths)).split(
              depths
          )
      ]
      stages = []
      for stage_idx in range(num_stages):
          out_chs = embed_dims[stage_idx]
          stage = DaVitStage(
              in_chs,
              out_chs,
              depth=depths[stage_idx],
              downsample=stage_idx > 0,
              attn_types=attn_types,
              num_heads=num_heads[stage_idx],
              window_size=window_size,
              mlp_ratio=mlp_ratio,
              qkv_bias=qkv_bias,
              drop_path_rates=dpr[stage_idx],
              norm_layer=norm_layer,
              norm_layer_cl=norm_layer_cl,
              ffn=ffn,
              cpe_act=cpe_act,
          )
          in_chs = out_chs
          stages.append(stage)
          self.feature_info += [
              dict(
                  num_chs=out_chs,
                  reduction=2,
                  module=f"stages.{stage_idx}",
              )
          ]
  
      self.stages = nn.Sequential(*stages)
  
      # if head_norm_first == true, norm -> global pool -> fc ordering, like most other nets
      # otherwise pool -> norm -> fc, the default DaViT order, similar to ConvNeXt
      # FIXME generalize this structure to ClassifierHead
      if head_norm_first:
          self.norm_pre = norm_layer(self.num_features)
          self.head = ClassifierHead(
              self.num_features,
              num_classes,
              pool_type=global_pool,
              drop_rate=self.drop_rate,
          )
      else:
          self.norm_pre = nn.Identity()
          self.head = NormMlpClassifierHead(
              self.num_features,
              num_classes,
              pool_type=global_pool,
              drop_rate=self.drop_rate,
              norm_layer=norm_layer,
          )
      self.apply(self._init_weights)
  
imports: from functools import partial
  from typing import Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import DropPath, to_2tuple, trunc_normal_, Mlp, LayerNorm2d, get_norm_layer, use_fused_attn
  from timm.layers import NormMlpClassifierHead, ClassifierHead
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._manipulate import checkpoint_seq
  from ._registry import generate_default_cfgs, register_model
label_desc: Use image node_id 1 to create a new DaVit object from inherited base classes: nn with example: obj = DaVit(in_chans, depths, embed_dims, num_heads, window_size, mlp_ratio, qkv_bias, norm_layer, norm_layer_cl, norm_eps, attn_types, ffn, cpe_act, drop_rate, drop_path_rate, num_classes, global_pool, head_norm_first)