Args: x (tensor): input features with shape of (num_windows*B, N, C) mask (tensor | None, Optional): mask with shape of (num_windows, Wh*Ww, Wh*Ww), value should be between (-inf, 0].
(self, x, mask=None)
| 76 | trunc_normal_init(self.relative_position_bias_table, std=0.02) |
| 77 | |
| 78 | def forward(self, x, mask=None): |
| 79 | """ |
| 80 | Args: |
| 81 | |
| 82 | x (tensor): input features with shape of (num_windows*B, N, C) |
| 83 | mask (tensor | None, Optional): mask with shape of (num_windows, |
| 84 | Wh*Ww, Wh*Ww), value should be between (-inf, 0]. |
| 85 | """ |
| 86 | B, N, C = x.shape |
| 87 | qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, |
| 88 | C // self.num_heads).permute(2, 0, 3, 1, 4) |
| 89 | # make torchscript happy (cannot use tensor as tuple) |
| 90 | q, k, v = qkv[0], qkv[1], qkv[2] |
| 91 | |
| 92 | q = q * self.scale |
| 93 | attn = (q @ k.transpose(-2, -1)) |
| 94 | |
| 95 | relative_position_bias = self.relative_position_bias_table[ |
| 96 | self.relative_position_index.view(-1)].view( |
| 97 | self.window_size[0] * self.window_size[1], |
| 98 | self.window_size[0] * self.window_size[1], |
| 99 | -1) # Wh*Ww,Wh*Ww,nH |
| 100 | relative_position_bias = relative_position_bias.permute( |
| 101 | 2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww |
| 102 | attn = attn + relative_position_bias.unsqueeze(0) |
| 103 | |
| 104 | if mask is not None: |
| 105 | nW = mask.shape[0] |
| 106 | attn = attn.view(B // nW, nW, self.num_heads, N, |
| 107 | N) + mask.unsqueeze(1).unsqueeze(0) |
| 108 | attn = attn.view(-1, self.num_heads, N, N) |
| 109 | attn = self.softmax(attn) |
| 110 | |
| 111 | attn = self.attn_drop(attn) |
| 112 | |
| 113 | x = (attn @ v).transpose(1, 2).reshape(B, N, C) |
| 114 | x = self.proj(x) |
| 115 | x = self.proj_drop(x) |
| 116 | return x |
| 117 | |
| 118 | @staticmethod |
| 119 | def double_step_seq(step1, len1, step2, len2): |
nothing calls this directly
no outgoing calls
no test coverage detected