(self, query, hw_shape)
| 175 | self.drop = build_dropout(dropout_layer) |
| 176 | |
| 177 | def forward(self, query, hw_shape): |
| 178 | B, L, C = query.shape |
| 179 | H, W = hw_shape |
| 180 | assert L == H * W, 'input feature has wrong size' |
| 181 | query = query.view(B, H, W, C) |
| 182 | |
| 183 | # pad feature maps to multiples of window size |
| 184 | pad_r = (self.window_size - W % self.window_size) % self.window_size |
| 185 | pad_b = (self.window_size - H % self.window_size) % self.window_size |
| 186 | query = F.pad(query, (0, 0, 0, pad_r, 0, pad_b)) |
| 187 | H_pad, W_pad = query.shape[1], query.shape[2] |
| 188 | |
| 189 | # cyclic shift |
| 190 | if self.shift_size > 0: |
| 191 | shifted_query = torch.roll( |
| 192 | query, |
| 193 | shifts=(-self.shift_size, -self.shift_size), |
| 194 | dims=(1, 2)) |
| 195 | |
| 196 | # calculate attention mask for SW-MSA |
| 197 | img_mask = torch.zeros((1, H_pad, W_pad, 1), device=query.device) |
| 198 | h_slices = (slice(0, -self.window_size), |
| 199 | slice(-self.window_size, |
| 200 | -self.shift_size), slice(-self.shift_size, None)) |
| 201 | w_slices = (slice(0, -self.window_size), |
| 202 | slice(-self.window_size, |
| 203 | -self.shift_size), slice(-self.shift_size, None)) |
| 204 | cnt = 0 |
| 205 | for h in h_slices: |
| 206 | for w in w_slices: |
| 207 | img_mask[:, h, w, :] = cnt |
| 208 | cnt += 1 |
| 209 | |
| 210 | # nW, window_size, window_size, 1 |
| 211 | mask_windows = self.window_partition(img_mask) |
| 212 | mask_windows = mask_windows.view( |
| 213 | -1, self.window_size * self.window_size) |
| 214 | attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) |
| 215 | attn_mask = attn_mask.masked_fill(attn_mask != 0, |
| 216 | float(-100.0)).masked_fill( |
| 217 | attn_mask == 0, float(0.0)) |
| 218 | else: |
| 219 | shifted_query = query |
| 220 | attn_mask = None |
| 221 | |
| 222 | # nW*B, window_size, window_size, C |
| 223 | query_windows = self.window_partition(shifted_query) |
| 224 | # nW*B, window_size*window_size, C |
| 225 | query_windows = query_windows.view(-1, self.window_size**2, C) |
| 226 | |
| 227 | # W-MSA/SW-MSA (nW*B, window_size*window_size, C) |
| 228 | attn_windows = self.w_msa(query_windows, mask=attn_mask) |
| 229 | |
| 230 | # merge windows |
| 231 | attn_windows = attn_windows.view(-1, self.window_size, |
| 232 | self.window_size, C) |
| 233 | |
| 234 | # B H' W' C |
nothing calls this directly
no test coverage detected