B, N, C = x.shape
# self.qkv生成qkv output of self.qkv: [5640, 49, 288]
qkv_x = self.qkv(x).reshape(B, N, 3, self.num_heads,
C // self.num_heads).permute(2, 0, 3, 1, 4)
qkv_y = self.qkv(y).reshape(B, N, 3, self.num_heads,
C // self.num_heads).permute(2, 0, 3, 1, 4)
# qkv.shape = [3, 5640, 3, 49, 32]
# make torchscript happy (cannot use tensor as tuple)
qx, kx, vx = qkv_x[0], qkv_x[1], qkv_x[2]
qy, ky, vy = qkv_y[0], qkv_y[1], qkv_y[2]
qx = qx * self.scale
# qy = qy * self.scale
attn = (qx @ ky.transpose(-2, -1))
# 定义一个相对位置的表:
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) # Wh*Ww,Wh*Ww,nH
relative_position_bias = relative_position_bias.permute(
2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
attn = attn + relative_position_bias.unsqueeze(0)
if mask is not None:
nW = mask.shape[0]
attn = attn.view(B // nW, nW, self.num_heads, N,
N) + mask.unsqueeze(1).unsqueeze(0)
attn = attn.view(-1, self.num_heads, N, N)
attn = self.softmax(attn)
attn = self.attn_drop(attn)
x = (attn @ vx).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
y = (attn @ vy).transpose(1, 2).reshape(B, N, C)
y = self.proj(y)
y = self.proj_drop(y)
return x, y