class_docstr: Multi-head attention layer.
  
  Args:
      dim (`int`):
          number of features
      num_heads (`int=8`):
          number of heads
      qkv_bias (`bool=False`):
          if add bias to qkv layer
      qk_scale (`float=None`):
          number to scale qk
      attn_drop_ratio (`float=0.`):
          drop rate of attention layer
      proj_drop_ratio (`float=0.`):
          drop rate of projection layer
      with_qkv (`bool=True`):
          if use qkv layer
  
  Example:
      >>> import torch
      >>> from towhee.models.layers.attention import MultiHeadAttention
      >>>
      >>> test_shape = (1, 196+1, 768)  # shape of output from patch_embed
      >>> input_x = torch.rand(test_shape)
      >>> model = MultiHeadAttention(dim=test_shape[2])
      >>> out = model.forward(input_x)
      >>> print(out.shape)
      torch.Size([1, 197, 768])
code: def __init__(
      self,
      dim,
      num_heads=8,
      qkv_bias=False,
      qk_scale=None,
      attn_drop_ratio=0.0,
      proj_drop_ratio=0.0,
      with_qkv=True,
  ):
      super().__init__()
      _ = with_qkv
      self.num_heads = num_heads
      self.head_dim = dim // num_heads
      self.scale = qk_scale or self.head_dim**-0.5
  
      # A = Q*K^T
      self.matmul1 = Einsum("bhid,bhjd->bhij")
      # attn = A*V
      self.matmul2 = Einsum("bhij,bhjd->bhid")
  
      self.qkv = Linear(
          dim, dim * 3, qkv_bias
      )  # pylint: disable=too-many-function-args
      self.attn_drop = Dropout(
          attn_drop_ratio
      )  # pylint: disable=too-many-function-args
      self.proj = Linear(
          dim, dim
      )  # pylint: disable=too-many-function-args
      self.proj_drop = Dropout(
          proj_drop_ratio
      )  # pylint: disable=too-many-function-args
      self.softmax = Softmax(
          dim=-1
      )  # pylint: disable=unexpected-keyword-arg
  
      self.attn_cam = None
      self.attn = None
      self.v = None
      self.v_cam = None
      self.attn_gradients = None
  
imports: import os
  import math
  import numpy
  import torch
  from torch import nn
  from towhee.models.layers.layers_with_relprop import Einsum, Linear, Dropout, Softmax
label_desc: Use image node_id 1 to create a new MultiHeadAttention object from inherited base classes: nn with example: obj = MultiHeadAttention(dim, num_heads, qkv_bias, qk_scale, attn_drop_ratio, proj_drop_ratio, with_qkv)