class_docstr: A multiscale attention block.
  compare to a conventional attention block, a multiscale attention block optionally
  supports pooling (either before or after qkv projection). If pooling is not used, a
  multiscale attention block is equivalent to a conventional attention block.
  ::
                                 Input
                                   |
                  |----------------|-----------------|
                  ↓                ↓                 ↓
                Linear           Linear            Linear
                  &                &                 &
               Pool (Q)         Pool (K)          Pool (V)
                  → -------------- ←                 |
                           ↓                         |
                     MatMul & Scale                  |
                           ↓                         |
                        Softmax                      |
                           → ----------------------- ←
                                       ↓
                                 MatMul & Scale
                                       ↓
                                    DropOut
  Args:
      dim(`int`):
          Input feature dimension.
      num_heads(`int`):
          number of heads in the attention layer.
      qkv_bias(`bool`):
          If set to False, the qkv layer will not learn an additive bias.
      dropout_rate(`float`):
          Dropout rate.
      kernel_q(`_size_3_t`):
          Pooling kernel size for q. If both pooling kernel
          size and pooling stride size are 1 for all the dimensions, pooling is disabled.
      kernel_kv(`_size_3_t`):
          Pooling kernel size for kv. If both pooling kernel size and pooling stride size
          are 1 for all the dimensions, pooling is disabled.
      stride_q(`_size_3_t`):
          Pooling kernel stride for q.
      stride_kv(`_size_3_t`):
          Pooling kernel stride for kv.
      norm_layer(`nn.Module`):
          normalization layer used after pooling.
      has_cls_embed(`bool`):
          If set to True, the first token of the input tensor
          should be a cls token. Otherwise, the input tensor does not contain a cls token.
          Pooling is not applied to the cls token.
      pool_mode(`str`):
          Pooling mode. Option includes "conv" (learned pooling), "avg"
          (average pooling), and "max" (max pooling).
      pool_first(`bool`):
          If set to True, pool is applied before qkv projection.
          Otherwise, pool is applied after qkv projection.
class_docstr_tok: ["A","multiscale","attention","block",".","compare","to","a","conventional","attention","block",",","a","multiscale","attention","block","optionally","supports","pooling","(","either","before","or","after","qkv","projection",")",".","If","pooling","is","not","used",",","a","multiscale","attention","block","is","equivalent","to","a","conventional","attention","block",".",":",":","Input","|","|","--","--","--","--","--","--","--","--","|","--","--","--","--","--","--","--","--","-|","\u2193","\u2193","\u2193","Linear","Linear","Linear","&","&","&","Pool","(","Q",")","Pool","(","K",")","Pool","(","V",")","\u2192","--","--","--","--","--","--","--","\u2190","|","\u2193","|","MatMul","&","Scale","|","\u2193","|","Softmax","|","\u2192","--","--","--","--","--","--","--","--","--","--","--","-","\u2190","\u2193","MatMul","&","Scale","\u2193","DropOut","Args",":","dim","(","`","int","`",")",":","Input","feature","dimension",".","num_heads","(","`","int","`",")",":","number","of","heads","in","the","attention","layer",".","qkv_bias","(","`","bool","`",")",":","If","set","to","False",",","the","qkv","layer","will","not","learn","an","additive","bias",".","dropout_rate","(","`","float","`",")",":","Dropout","rate",".","kernel_q","(","`","_size_3_t","`",")",":","Pooling","kernel","size","for","q",".","If","both","pooling","kernel","size","and","pooling","stride","size","are","1","for","all","the","dimensions",",","pooling","is","disabled",".","kernel_kv","(","`","_size_3_t","`",")",":","Pooling","kernel","size","for","kv",".","If","both","pooling","kernel","size","and","pooling","stride","size","are","1","for","all","the","dimensions",",","pooling","is","disabled",".","stride_q","(","`","_size_3_t","`",")",":","Pooling","kernel","stride","for","q.","stride_kv","(","`","_size_3_t","`",")",":","Pooling","kernel","stride","for","kv",".","norm_layer","(","`","nn.Module","`",")",":","normalization","layer","used","after","pooling",".","has_cls_embed","(","`","bool","`",")",":","If","set","to","True",",","the","first","token","of","the","input","tensor","should","be","a","cls","token",".","Otherwise",",","the","input","tensor","does","not","contain","a","cls","token",".","Pooling","is","not","applied","to","the","cls","token",".","pool_mode","(","`","str","`",")",":","Pooling","mode",".","Option","includes","``","conv","''","(","learned","pooling",")",",","``","avg","''","(","average","pooling",")",",","and","``","max","''","(","max","pooling",")",".","pool_first","(","`","bool","`",")",":","If","set","to","True",",","pool","is","applied","before","qkv","projection",".","Otherwise",",","pool","is","applied","after","qkv","projection","."]
code: def __init__(
      self,
      dim,
      num_heads=8,
      qkv_bias=False,
      dropout_rate=0.0,
      kernel_q=(1, 1, 1),
      kernel_kv=(1, 1, 1),
      stride_q=(1, 1, 1),
      stride_kv=(1, 1, 1),
      norm_layer=nn.LayerNorm,
      has_cls_embed=True,
      pool_mode=nn.Conv3d,
      pool_first=False,
  ) -> None:
      super().__init__()
      self.dim = dim
      self.num_heads = num_heads
      self.qkv_bias = qkv_bias
      self.dropout_rate = dropout_rate
      self.kernel_q = kernel_q
      self.kernel_kv = kernel_kv
      self.stride_q = stride_q
      self.stride_kv = stride_kv
      self.norm_layer = norm_layer
      self.has_cls_embed = has_cls_embed
      self.pool_mode = pool_mode
      self.pool_first = pool_first
  
      assert self.pool_mode in [nn.Conv3d, nn.AvgPool3d, nn.MaxPool3d]
  
      self.head_dim = self.dim // self.num_heads
      self.scale = self.head_dim**-0.5
      self.padding_q = [int(q // 2) for q in self.kernel_q]
      self.padding_kv = [int(kv // 2) for kv in self.kernel_kv]
  
      self.q = nn.Linear(self.dim, self.dim, bias=self.qkv_bias)
      self.k = nn.Linear(self.dim, self.dim, bias=self.qkv_bias)
      self.v = nn.Linear(self.dim, self.dim, bias=self.qkv_bias)
      self.proj = nn.Linear(self.dim, self.dim)
      if self.dropout_rate > 0.0:
          self.proj_drop = nn.Dropout(self.dropout_rate)
  
      # Skip pooling with kernel and stride size of (1, 1, 1).
      if (
          self.kernel_q is not None
          and numpy.prod(self.kernel_q) == 1
          and numpy.prod(self.stride_q) == 1
      ):
          self.kernel_q = None
      if (
          self.kernel_kv is not None
          and numpy.prod(self.kernel_kv) == 1
          and numpy.prod(self.stride_kv) == 1
      ):
          self.kernel_kv = None
  
      if self.pool_mode in (nn.AvgPool3d, nn.MaxPool3d):
          pool_op = (
              nn.MaxPool3d
              if pool_mode == nn.MaxPool3d
              else nn.AvgPool3d
          )
          self.pool_q = (
              pool_op(
                  self.kernel_q,
                  self.stride_q,
                  self.padding_q,
                  ceil_mode=False,
              )
              if self.kernel_q is not None
              else None
          )
          self.pool_k = (
              pool_op(
                  self.kernel_kv,
                  self.stride_kv,
                  self.padding_kv,
                  ceil_mode=False,
              )
              if self.kernel_kv is not None
              else None
          )
          self.pool_v = (
              pool_op(
                  self.kernel_kv,
                  self.stride_kv,
                  self.padding_kv,
                  ceil_mode=False,
              )
              if self.kernel_kv is not None
              else None
          )
      elif self.pool_mode == nn.Conv3d:
          self.pool_q = (
              nn.Conv3d(
                  self.head_dim,
                  self.head_dim,
                  self.kernel_q,
                  stride=self.stride_q,
                  padding=self.padding_q,
                  groups=self.head_dim,
                  bias=False,
              )
              if self.kernel_q is not None
              else None
          )
          self.norm_q = (
              self.norm_layer(self.head_dim)
              if self.kernel_q is not None
              else None
          )
          self.pool_k = (
              nn.Conv3d(
                  self.head_dim,
                  self.head_dim,
                  self.kernel_kv,
                  stride=self.stride_kv,
                  padding=self.padding_kv,
                  groups=self.head_dim,
                  bias=False,
              )
              if self.kernel_kv is not None
              else None
          )
          self.norm_k = (
              self.norm_layer(self.head_dim)
              if self.kernel_kv is not None
              else None
          )
          self.pool_v = (
              nn.Conv3d(
                  self.head_dim,
                  self.head_dim,
                  self.kernel_kv,
                  stride=self.stride_kv,
                  padding=self.padding_kv,
                  groups=self.head_dim,
                  bias=False,
              )
              if self.kernel_kv is not None
              else None
          )
          self.norm_v = (
              self.norm_layer(self.head_dim)
              if self.kernel_kv is not None
              else None
          )
      else:
          raise NotImplementedError("Unsupported model.")
  
file_path: towhee/towhee/models/layers/multi_scale_attention.py
imports: import torch
  from torch import nn
  from typing import List, Tuple
  import numpy
  from towhee.models.layers.pool_attention import AttentionPool
label_desc: Use image node_id 1 to create a new MultiScaleAttention object from inherited base classes: nn with example: obj = MultiScaleAttention(dim, num_heads, qkv_bias, dropout_rate, kernel_q, kernel_kv, stride_q, stride_kv, norm_layer, has_cls_embed, pool_mode, pool_first)