class_docstr: Modified Attention.
  
  Args:
      dim (`int`): feature dimension
      num_heads (`int`): number of attention heads
      qkv_bias (`bool`): flag to use bias for qkv
      qk_scale (`int`): Number to scale qk
      attn_drop (`float`): drop rate of attention
      proj_drop (`float`): drop rate of projection
      sr_ratio (`int`): shunt rate
class_docstr_tok: ["Modified","Attention",".","Args",":","dim","(","`","int","`",")",":","feature","dimension","num_heads","(","`","int","`",")",":","number","of","attention","heads","qkv_bias","(","`","bool","`",")",":","flag","to","use","bias","for","qkv","qk_scale","(","`","int","`",")",":","Number","to","scale","qk","attn_drop","(","`","float","`",")",":","drop","rate","of","attention","proj_drop","(","`","float","`",")",":","drop","rate","of","projection","sr_ratio","(","`","int","`",")",":","shunt","rate"]
code: def forward(self, x, h, w):
      b, n, c = x.shape
      q = (
          self.q(x)
          .reshape(b, n, self.num_heads, c // self.num_heads)
          .permute(0, 2, 1, 3)
      )
      if self.sr_ratio > 1:
          x_ = x.permute(0, 2, 1).reshape(b, c, h, w)
          x_1 = self.act(
              self.norm1(
                  self.sr1(x_).reshape(b, c, -1).permute(0, 2, 1)
              )
          )
          x_2 = self.act(
              self.norm2(
                  self.sr2(x_).reshape(b, c, -1).permute(0, 2, 1)
              )
          )
          kv1 = (
              self.kv1(x_1)
              .reshape(
                  b, -1, 2, self.num_heads // 2, c // self.num_heads
              )
              .permute(2, 0, 3, 1, 4)
          )
          kv2 = (
              self.kv2(x_2)
              .reshape(
                  b, -1, 2, self.num_heads // 2, c // self.num_heads
              )
              .permute(2, 0, 3, 1, 4)
          )
          k1, v1 = kv1[0], kv1[1]
          k2, v2 = kv2[0], kv2[1]
          attn1 = (
              q[:, : self.num_heads // 2] @ k1.transpose(-2, -1)
          ) * self.scale
          attn1 = attn1.softmax(dim=-1)
          attn1 = self.attn_drop(attn1)
          v1 = v1 + self.local_conv1(
              v1.transpose(1, 2)
              .reshape(b, -1, c // 2)
              .transpose(1, 2)
              .view(b, c // 2, h // self.sr_ratio, w // self.sr_ratio)
          ).view(b, c // 2, -1).view(
              b, self.num_heads // 2, c // self.num_heads, -1
          ).transpose(
              -1, -2
          )
          x1 = (attn1 @ v1).transpose(1, 2).reshape(b, n, c // 2)
          attn2 = (
              q[:, self.num_heads // 2 :] @ k2.transpose(-2, -1)
          ) * self.scale
          attn2 = attn2.softmax(dim=-1)
          attn2 = self.attn_drop(attn2)
          v2 = v2 + self.local_conv2(
              v2.transpose(1, 2)
              .reshape(b, -1, c // 2)
              .transpose(1, 2)
              .view(
                  b,
                  c // 2,
                  h * 2 // self.sr_ratio,
                  w * 2 // self.sr_ratio,
              )
          ).view(b, c // 2, -1).view(
              b, self.num_heads // 2, c // self.num_heads, -1
          ).transpose(
              -1, -2
          )
          x2 = (attn2 @ v2).transpose(1, 2).reshape(b, n, c // 2)
  
          x = torch.cat([x1, x2], dim=-1)
      else:
          kv = (
              self.kv(x)
              .reshape(b, -1, 2, self.num_heads, c // self.num_heads)
              .permute(2, 0, 3, 1, 4)
          )
          k, v = kv[0], kv[1]
  
          attn = (q @ k.transpose(-2, -1)) * self.scale
          attn = attn.softmax(dim=-1)
          attn = self.attn_drop(attn)
  
          x = (attn @ v).transpose(1, 2).reshape(
              b, n, c
          ) + self.local_conv(
              v.transpose(1, 2)
              .reshape(b, n, c)
              .transpose(1, 2)
              .view(b, c, h, w)
          ).view(
              b, c, n
          ).transpose(
              1, 2
          )
      x = self.proj(x)
      x = self.proj_drop(x)
  
      return x
  
imports: import torch
  from torch import nn
  from towhee.models.layers.patch_embed2d import PatchEmbed2D
  from towhee.models.layers.mlp import Mlp
  from towhee.models.layers.droppath import DropPath
  from towhee.models.utils.init_vit_weights import init_vit_weights
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj.forward(x, h, w) and returns: x