Swing transformer details-2

chart 1 2 individual MSA. Corresponding to the following depth Medium 2.

 

chart 2 SW-MSA technological process

 

1. cyclic shift + reverse cyclic shift

chart 3 cyclic shift, shift_size = 7 // 2 = 3  

 

depth = [2, 2, 6, 2] # MSA The number of 
[SwinTransformerBlock(...,shift_size=0 if (i % 2 == 0) else window_size // 2,...)
for i in range(depth)] # window_size

therefore , In the diagram above , W-MSA Of shift_size = 0,SW-MSA in mask=self.attn_mask,  shift_size = 7 // 2 = 3.

window_size = 7
shift_size = 7 // 2
''' Construct multidimensional tensor '''
# x=np.arange(14*14*4*96).reshape(1,14,14,96*4)
x=np.arange(14*14).reshape(14,14)
x=torch.from_numpy(x)
print(x)
if shift_size > 0:
    shifted_x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(0, 1))
    #shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
    print("---------cyclic shift---------")
else:
    shifted_x = x
print(shifted_x)  

# reverse cyclic shift
if shift_size > 0:
    x = torch.roll(shifted_x, shifts=(shift_size, shift_size), dims=(0, 1))
    print("---------reverse cyclic shift---------")
    print(x)  
else:
    x = shifted_x

chart 4 SW Specific changes

2. SW-MSA

shift_size = 3
window_size = 7
if shift_size > 0:
    input_resolution = (14, 14)
    # calculate attention mask for SW-MSA
    H, W = input_resolution
    # img_mask = torch.zeros((H, W))  # H W
    img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
    h_slices = (slice(0, -window_size),
                slice(-window_size, -shift_size),
                slice(-shift_size, None))
    w_slices = (slice(0, -window_size),
                slice(-window_size, -shift_size),
                slice(-shift_size, None))
    cnt = 0
    for h in h_slices:
        for w in w_slices:
            # img_mask[h, w] = cnt
            img_mask[:, h, w, :] = cnt
            cnt += 1
            
    mask_windows = window_partition(img_mask, window_size)  # nW, window_size, window_size ,  here  nW = 4

    outputs = mask_windows.view(-1, window_size, window_size)
    outputs_1 = torch.stack((outputs[0], outputs[1]), dim=1).view(-1, window_size, window_size*2)
    outputs_2 = torch.stack((outputs[2], outputs[3]), dim=1).view(-1, window_size, window_size*2)
    outputs = torch.stack((outputs_1, outputs_2), dim=1).view(-1, H, W)
    print(outputs)

    mask_windows = mask_windows.view(-1, window_size * window_size) # nW, window_size * window_size
    print(mask_windows)
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) # (nW, 1, window_size * window_size) - (nW, window_size * window_size, 1) 
                                                               #  radio broadcast  -> (nW, window_size * window_size, window_size * window_size) - (nW, window_size * window_size, window_size * window_size)
    print(attn_mask[1][0] == attn_mask[1][4])
    print(attn_mask[1][4].view(window_size, window_size))
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) #  take   Not  0  The replacement of is  -100, 0 Replace with  0.
    print(attn_mask.shape)

attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C

mask = attn_mask
if mask is not None:
    nW = mask.shape[0] #  How many are a picture divided into windows
    attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) # torch.Size([128, 4, 12, 49, 49]) torch.Size([1, 4, 1, 49, 49])  radio broadcast 
    attn = attn.view(-1, self.num_heads, N, N)
    attn = self.softmax(attn)  

Use attn + mask , Make the current position ( Like above [1][4]), Non adjacent areas in the image -100, Equivalent to calculating softmax Do not consider these areas .

3. W-MSA and MSA Complexity comparison of + 4.  Overall flow chart

          Please refer to ： Detailed explanation of the paper ：Swin Transformer - You know