class_docstr: Window based multi-head self attention (W-MSA) module with relative position bias.
  It supports both of shifted and non-shifted window.
  Args:
      dim (`int`):
          Number of input channels.
      window_size (`tuple[int]`):
          The height and width of the window.
      num_heads (`int`):
          Number of attention heads.
      qkv_bias (`bool`):
          If True, add a learnable bias to query, key, value. Default: True
      attn_drop (`float`):
          Dropout ratio of attention weight. Default: 0.0
      proj_drop (`float`):
          Dropout ratio of output. Default: 0.0
      pretrained_window_size (`tuple[int]`):
          The height and width of the window in pre-training.
code: def __init__(
      self,
      dim,
      window_size,
      num_heads,
      qkv_bias=True,
      attn_drop=0.0,
      proj_drop=0.0,
      is_v2=False,
      pretrained_window_size=None,
  ):
      super().__init__()
      if pretrained_window_size is None:
          pretrained_window_size = [0, 0]
      self.dim = dim
      self.window_size = window_size  # Wh, Ww
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = head_dim**-0.5
      self.is_v2 = is_v2
      self.pretrained_window_size = pretrained_window_size
  
      if self.is_v2:
          self.logit_scale = nn.Parameter(
              torch.log(10 * torch.ones((num_heads, 1, 1))),
              requires_grad=True,
          )
  
          # mlp to generate continuous relative position bias
          self.cpb_mlp = nn.Sequential(
              nn.Linear(2, 512, bias=True),
              nn.ReLU(inplace=True),
              nn.Linear(512, num_heads, bias=False),
          )
  
          # get relative_coords_table
          relative_coords_h = torch.arange(
              -(self.window_size[0] - 1),
              self.window_size[0],
              dtype=torch.float32,
          )
          relative_coords_w = torch.arange(
              -(self.window_size[1] - 1),
              self.window_size[1],
              dtype=torch.float32,
          )
          relative_coords_table = (
              torch.stack(
                  torch.meshgrid([relative_coords_h, relative_coords_w])
              )
              .permute(1, 2, 0)
              .contiguous()
              .unsqueeze(0)
          )  # 1, 2*Wh-1, 2*Ww-1, 2
          if pretrained_window_size[0] > 0:
              relative_coords_table[:, :, :, 0] /= (
                  pretrained_window_size[0] - 1
              )
              relative_coords_table[:, :, :, 1] /= (
                  pretrained_window_size[1] - 1
              )
          else:
              relative_coords_table[:, :, :, 0] /= (
                  self.window_size[0] - 1
              )
              relative_coords_table[:, :, :, 1] /= (
                  self.window_size[1] - 1
              )
          relative_coords_table *= 8  # normalize to -8, 8
          relative_coords_table = (
              torch.sign(relative_coords_table)
              * torch.log2(torch.abs(relative_coords_table) + 1.0)
              / np.log2(8)
          )
  
          self.register_buffer(
              "relative_coords_table", relative_coords_table
          )
          self.qkv = nn.Linear(dim, dim * 3, bias=False)
          if qkv_bias:
              self.q_bias = nn.Parameter(torch.zeros(dim))
              self.v_bias = nn.Parameter(torch.zeros(dim))
          else:
              self.q_bias = None
              self.v_bias = None
      else:
          # define a parameter table of relative position bias
          self.relative_position_bias_table = nn.Parameter(
              torch.zeros(
                  (2 * window_size[0] - 1) * (2 * window_size[1] - 1),
                  num_heads,
              )
          )  # 2*Wh-1 * 2*Ww-1, nH
          trunc_normal_(self.relative_position_bias_table, std=0.02)
          self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
      # get pair-wise relative position index for each token inside the window
      coords_h = torch.arange(self.window_size[0])
      coords_w = torch.arange(self.window_size[1])
      coords = torch.stack(
          torch.meshgrid([coords_h, coords_w])
      )  # 2, Wh, Ww
      coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
      relative_coords = (
          coords_flatten[:, :, None] - coords_flatten[:, None, :]
      )  # 2, Wh*Ww, Wh*Ww
      relative_coords = relative_coords.permute(
          1, 2, 0
      ).contiguous()  # Wh*Ww, Wh*Ww, 2
      relative_coords[:, :, 0] += (
          self.window_size[0] - 1
      )  # shift to start from 0
      relative_coords[:, :, 1] += self.window_size[1] - 1
      relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
      relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
      self.register_buffer(
          "relative_position_index", relative_position_index
      )
  
      self.attn_drop = nn.Dropout(attn_drop)
      self.proj = nn.Linear(dim, dim)
      self.proj_drop = nn.Dropout(proj_drop)
  
      self.softmax = nn.Softmax(dim=-1)
  
file_path: towhee/towhee/models/layers/window_attention.py
imports: import torch
  import torch.nn.functional
  from torch import nn
  import numpy
  from typing import Optional
  from towhee.models.utils.weight_init import trunc_normal_
label_desc: Use image node_id 1 to create a new WindowAttention object from inherited base classes: nn with example: obj = WindowAttention(dim, window_size, num_heads, qkv_bias, attn_drop, proj_drop, is_v2, pretrained_window_size)