if self.shift_size > 0: # calculate attention mask for SW-MSA H, W = self.input_resolution img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) else: attn_mask = None

这段代码是 ChitGPT 中的一部分，用于实现基于滑动窗口的多尺度自注意力机制（SW-MSA）。主要是计算用于掩盖不相关像素的注意力掩码。如果 shift_size 大于 0，就会生成一个大小为 H x W 的图像掩码，然后将其分成若干个大小为 window_size x window_size 的窗口。对于每对窗口，将它们的编号相减，并用 -100.0 填充非零元素的位置，用 0.0 填充零元素的位置，生成一个注意力掩码。如果 shift_size 等于 0，则不需要掩码。