Transformer--相对位置编码

最近使用较多，不同于绝对位置编码，相对位置编码用在每次计算attention矩阵时
atten = softmax（Q*K+B）
这里就以Swin Transformer为例，直接上代码

1.即插即用的相对位置偏移代码

该代码可单独运行，在计算attn（=q*k）时可以直接与attn相加

import torch
import torch.nn as nn
from timm.models.layers import trunc_normal_

window_size = [7,7]
num_heads = 8


relative_position_bias_table = nn.Parameter(torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
coords_h = torch.arange(window_size[0])
coords_w = torch.arange(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] += window_size[0] - 1  # shift to start from 0
relative_coords[:, :, 1] += window_size[1] - 1
relative_coords[:, :, 0] *= 2 * window_size[1] - 1
relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
trunc_normal_(relative_position_bias_table, std=.02)

relative_position_bias = relative_position_bias_table[relative_position_index.view(-1)].view(
            window_size[0] * window_size[1], window_size[0] * window_size[1], -1)

relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()

2.原理解释（结合代码）

1.这里以window_size=7为例

        # calculate attention mask for SW-MSA 输入特征图（200，304，96）：
        Hp = int(np.ceil(H / self.window_size)) * self.window_size          # 200-->203
        Wp = int(np.ceil(W / self.window_size)) * self.window_size        # 304-->308      改成 7的倍数
        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1       ( 1,203,308,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
这段我理解的是，从H和W中切（-7到-3）和（-3到-1）出来，相关位置因为是后续pad上去的，所以mask的位置是个较大的数值（最大到cnt=9）。其他位置是原特征图上有的点，因此mask都为0。

        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1          1276,7,7,1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)                       # 1276 , 49
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)                                 # 1276 , 49 , 49
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

2.计算RPB

维度为（169，2）
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=.02)   # 截断标准化RPB。 截断正态分布随机数，均值mean,标准差stddev,不过只保留[mean-2stddev,mean+2stddev]范围内的随机数

3.计算索引index

# 初始化均为0的可学习变量，由于位置编码取值为【-6，6】共13个数，3个head
       # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(self.window_size[0])                     # [0, 1, 2, 3, 4, 5, 6]
        coords_w = torch.arange(self.window_size[1])                    # [0, 1, 2, 3, 4, 5, 6]
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww            (2,7,7)
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww           # (2,49)
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww    (2,49,49)
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += self.window_size[0] - 1  #  值加上6，变成从0开始。原理类似图中加2
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1        # x坐标乘13，最后再和y坐标加在一起
        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww        (49,49,2) -> (49,49)

4.最终计算attn矩阵

        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  
            #   (169,3) [49,49] --> ([49, 49, 3])
            # 这里我也有点疑虑，为啥（169，3）变成了（49，49，3）
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # （3，49，49）
        attn = attn + relative_position_bias.unsqueeze(0)