SG Former论文学习笔记-编程知识

超越SWin和CSWin Transformer的新模型

代码地址：https://github.com/OliverRensu/SG-Former

论文地址：https://arxiv.org/pdf/2308.12216.pdf

ViT在各种视觉任务中虽然成功，但它的计算成本随着Token序列长度的增加呈二次增长，这在处理大型特征图时大大限制了其性能。为了减轻计算成本，先前的研究依赖于要么限制在局部小区域内的细粒度Self-Attention，要么采用全局Self-Attention，但要缩短序列长度，从而导致粗粒度的问题。

Self-Guided Transformer（SG-Former），以实现具有自适应细粒度的有效全局Self-Attention。本文方法的核心思想是利用一个显著性图，通过混合尺度的Self-Attention估计并在训练过程中自我演化，以根据每个区域的重要性重新分配Token。直观地说，为显著区域分配更多的Token，以实现细粒度的注意力，同时将更少的Token分配给次要区域，以换取效率和全局感知字段。SG-Former在分类，检测和分割任务中都较最先进的模型有提升。

为了在高分辨率特征上计算Self-Attention，一些方法提出了将Self-Attention区域限制在局部窗口而不是整个特征图（即细粒度的局部Self-Attention）。例如，Swin Transformer设计了窗口注意力，而CSWin设计了交叉形状注意力。因此，这些方法牺牲了在每个Self-Attention层中建模全局信息的能力。另一流方法旨在在整个键-值特征映射中聚合Token，以减少全局序列长度（即粗粒度的全局注意力）。例如，金字塔视觉Transformer（PVT）使用大步幅的大核心来均匀聚合整个特征映射上的Token，从而导致整个特征映射上的均匀粗糙信息。

本文的Self-Guided Transformer（SG-Former），通过一种不断演化的Self-Attention设计，实现了具有自适应细粒度的全局关注。SG-Former的核心思想在于，其保留了整个特征图上的远程依赖性，同时根据图像区域的重要性重新分配Token。

也就是说模型会为显著区域分配更多的Token，以便每个Token可以与显著区域以细粒度进行交互，同时为了效率将更少的Token分配给次要区域。SG-Former以高效的全局感知场计算Self-Attention，同时自适应地关注显著区域中的细粒度信息。

如图2所示，SG-Former根据从自身获得的注意力图重新分配Token，例如在狗这样的显著区域分配更多的Token，而在墙这样的次要区域分配更少的Token。 PVT采用了预定义的策略均匀聚合Token。

具体来说，保留了 Query Token，但重新分配了Key和Value Token以实现高效的全局Self-Attention。图像区域的重要性，以得分图的形式，通过混合尺度的Self-Attention本身估计，并进一步用于引导Token的重新分配。

也就是说，给定输入图像，Token的重新分配是通过Self-Guided完成的，这意味着每个图像都经历了一个仅适用于自身的独特的Token重新分配。因此，重新分配的Token受到人类先验的影响较小。

此外，这种Self-Guided会随着训练过程中逐渐精确的注意力图预测而不断演化。注意力图极大地影响了重新分配的有效性，因此提出了一种混合尺度的Self-Attention，该Self-Attention在同一层中具有与Swin相同成本的各种粒度信息。混合尺度Self-Attention通过分组Head和为不同的注意力粒度多样化每个组来实现各种粒度信息。混合尺度Self-Attention还向整个Transformer提供混合尺度信息。

本文有以下贡献：

1、细粒度的局部和全局粗粒度信息在一个Self-Attention层内，通过统一的混合尺度信息提取。具有统一的局部全局混合尺度信息，预测注意力图以识别区域重要性。

2、有了注意力图，可以模拟Self-Guided Attention，自动定位显著区域，并使显著区域提取细粒度信息，而使次要区域提取粗粒度信息以减少计算成本。

3、在分类、目标检测和分割任务中相较最先进的模型有明显提升。

下面结合代码来对论文进行讲解，代码部分以sgformer_s的配置进行分解。

具体的参数配置可见下表

1.1、Overview

SG-Former的总体流程如图3所示。SG-Former与以前的CNN和Transformer模型共享相同的Patch嵌入层和四阶金字塔架构。

首先，图像 $X\in R^{H\times W\times 3}$ 通过输入级别的Patch嵌入层进行4倍降采样。两个阶段之间存在一个2倍速率的降采样层。因此，第 $i$ 个阶段的特征图的大小为 $\frac{H}{2^{i+1}} \times \frac{W}{2^{i+1}}$ 。除了最后一个阶段，每个阶段都有 $N_{i}$ 个Transformer Block，由两种类型的块重复组成：

混合尺度 Transformer Block
Self-Guided Transformer Block。

混合尺度Self-Attention提取混合尺度对象和多粒度信息，引导区域重要性。Self-Guided Self-Attention模型全局信息，同时根据混合尺度 Transformer Block的重要性信息保持显著区域的细粒度；

先来看看SG Former的代码总体结构，代码示例部分省略了部分参数初始化操作。

class SGFormer(nn.Module):def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dims=[64, 128, 256, 512],num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], num_stages=4, linear=False):super().__init__()self.num_classes = num_classesself.depths = depthsself.num_stages = num_stagesself.num_patches = img_size//4dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rulecur = 0for i in range(num_stages):if i == 0:patch_embed = Head(embed_dims[0])  #else:patch_embed = PatchMerging(dim=embed_dims[i - 1],out_dim=embed_dims[i])block = nn.ModuleList([Block(dim=embed_dims[i], mask=True if (j%2==1 and i<num_stages-1) else False, num_heads=num_heads[i], mlp_ratio=mlp_ratios[i], qkv_bias=qkv_bias, qk_scale=qk_scale,drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + j], norm_layer=norm_layer,sr_ratio=sr_ratios[i], linear=linear)for j in range(depths[i])])norm = norm_layer(embed_dims[i])cur += depths[i]setattr(self, f"patch_embed{i + 1}", patch_embed)setattr(self, f"block{i + 1}", block)setattr(self, f"norm{i + 1}", norm)self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches*self.num_patches, embed_dims[0]))  # fixed sin-cos embedding# classification headself.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity()self.apply(self._init_weights)def forward_features(self, x):B = x.shape[0]mask=Nonefor i in range(self.num_stages):patch_embed = getattr(self, f"patch_embed{i + 1}")block = getattr(self, f"block{i + 1}")norm = getattr(self, f"norm{i + 1}")x, H, W = patch_embed(x)  # [N 3136 64] #[N 784 128] #[N 196 256] #[N 49 512]if i==0:x+=self.pos_embed  # [1 3136 64]for blk in block:x, mask = blk(x, H, W, mask)x = norm(x) # [N 3136 64] [N 784 128] [N 196 256] [N 49 512]if i != self.num_stages - 1:x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous() # [N 64 56 56] [N 128 28 28] [N 256 14 14]return x.mean(dim=1)def forward(self, x):x = self.forward_features(x) # entryx = self.head(x)return x

图像 $X\in R^{H\times W\times 3}$ 首先经过Patch embedding(代码中为Head)

class Head(nn.Module):def __init__(self, n):super(Head, self).__init__()self.conv = nn.Sequential(Conv2d_BN(3, n, 3, 2, 1),nn.GELU(),Conv2d_BN(n, n, 3, 1, 1),nn.GELU(),Conv2d_BN(n, n, 3, 2, 1),)self.norm = nn.LayerNorm(n)self.apply(self._init_weights)def forward(self, x):x = self.conv(x)_, _, H, W = x.shape  # [N 64 56 56]x = x.flatten(2)#.transpose(1, 2)  # [N 64 3136]x = x.transpose(1, 2)x = self.norm(x) # [N 3136 64]return x, H,W

其实就是经过三个2d卷积加BN，中间插入两个GELU激活函数，sequential头尾的两个2d卷积stride为2进行降采样，将输入X=[N,3,224,224]进行4倍降采样为[N,64,56,56]，之后再进行维度转换加LayerNorm。

第一层会对Patch embedding的输出加上pos embedding(shape为[1,3136,64])。

中间的Down sampling embedding 则是由PatchMerging完成：

class PatchMerging(nn.Module):def __init__(self, dim, out_dim):super().__init__()self.dim = dimself.out_dim = out_dimself.act = nn.GELU()self.conv1 = Conv2d_BN(dim, out_dim, 1, 1, 0)self.conv2 = Conv2d_BN(out_dim, out_dim, 3, 2, 1, groups=out_dim)self.conv3 = Conv2d_BN(out_dim, out_dim, 1, 1, 0)def forward(self, x):# x B C H Wx = self.conv1(x)x = self.act(x)x = self.conv2(x)x = self.act(x)x = self.conv3(x)_, _, H, W = x.shapex = x.flatten(2).transpose(1, 2)return x, H, W

主要是起到一个降采样的功能。

接下来就是本文的重点模块Transformer Block，由两种类型的块重复组成：

混合尺度 Transformer Block
Self-Guided Transformer Block。

这两个模块包括基本的Transformer Block都由Block函数构建

class Block(nn.Module):def __init__(self, dim, mask, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1, linear=False):super().__init__()self.norm1 = norm_layer(dim)self.attn = Attention(dim, mask,num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio, linear=linear)# NOTE: drop path for stochastic depth, we shall see if this is better than dropout hereself.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()self.norm2 = norm_layer(dim)mlp_hidden_dim = int(dim * mlp_ratio)self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)self.apply(self._init_weights)def forward(self, x, H, W, mask):x_, mask = self.attn(self.norm1(x), H, W, mask)  # x[N 3136 64] mask[[N 3136],[N 3136]]x = x + self.drop_path(x_)x = x + self.drop_path(self.mlp(self.norm2(x), H, W))return x, mask

Block函数中最重要的就是Attention

1.2、 Transformer Block

通过两种自注意机制，相应地设计了两种类型的Transformer Block。这两种Transformer Block只在注意层上有所不同，而其他所有方面都保持相同：

如图3所示，前3个阶段使用我们提出的混合尺度或Self-Guided Transformer Block进行定制，而对于最后一个阶段，使用了一个标准的 Transformer Block，这是根据以前的Transformer 进行的。请注意，前3个阶段（即N1、N2和N3）的 Transformer Block数量是偶数，而最后一个阶段（即N4）可以是偶数或奇数。

1.3、Hybrid-scale Attention

混合尺度注意力有2个目的：

提取混合尺度的全局和细粒度信息，而不比Swin Transformer中的窗口注意力多的计算成本
为Self-Guided 注意力提供重要性

如图5所示，输入特征X被投影到 Query (Q)、Key(K)和Value(V)。然后，Multi-Head Self-Attention采用H个独立Head。通常，这些H个独立Head在相同的局部区域内执行操作，因此缺乏Head的多样性。

相比之下，本文将H个Head均匀分成h个组，并将混合尺度和多感受野的注意力注入到这些h个组中，每个组中有 $\frac{H}{h}$ 个头（从代码上看，分了一半的head做global注意力，一半做local注意力）。在属于第 $j$ 组的第 $i$ 个Head中，使用尺度 $s_{j}$ （其中 $j=1,2,3...h$ ），将{K,V}中的每个 $s_{j} \times s_{j}$ Token 合并为一个 Token 。接下来，将{Q,K,V}分成窗口（和swin一样）。{K,V}的窗口大小设置为M，并在所有组中保持不变。为了使{Q}和{K,V}的窗口大小与{K,V}中的 Token 合并对齐，{Q}的窗口大小选择为 $s_{j}M\times s_{j}M$ ，是{K,V}的窗口大小的 $s_{j}$ 倍：

其中 $Merge(X,s_{j})$ 表示将中的每个 $s_{j} \times s_{j}$ Token 合并为一个 Token ，这是通过步长 $s_{j}$ 的卷积实现的。特殊情况是当 $s_{j}$ 等于1时，不进行 Token 合并，{Q,K,V}具有相同大小的窗口。

其中 $P(X,s_{j}M)$ 表示具有窗口大小 $s_{j}M\times s_{j}M$ 的窗口划分。 $Atten_{i}$ 是注意力图。有一个特殊情况： $s_{j}M\times s_{j}M$ 等于 $h\times w$ ，不需要窗口分割，{K,V}中的所有 Token 都由{Q}关注，从而实现全局信息提取。

Token 的重要性被视为所有 Token 和当前 Token 的乘积之和：

其中S是通过对所有 $S_{i}$ 求和得到的最终注意力图，用于混合尺度引导，提供全局和细粒度信息。对应这部分代码

# global
q1 = self.q1(x).reshape(B, N, self.num_heads//2, C // self.num_heads).permute(0, 2, 1, 3) # [N 1 3136 32] # [N 2 784 32] # [N 4 196 32]
x_ = x.permute(0, 2, 1).reshape(B, C, H, W) # [N 64 56 56] # [N 128 28 28] # [N 256 14 14]
x_1 = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) # [N 49 64] sr --> conv2d(64, 64) # [N 49 128] sr --> conv2d(128, 128) # [N 49 256] sr --> conv2d(256, 256)
x_1 = self.act(self.norm(x_1)) # [N 49 64] # [N 49 128] # [N 49 256]
kv1 = self.kv1(x_1).reshape(B, -1, 2, self.num_heads//2, C // self.num_heads).permute(2, 0, 3, 1, 4) # [2 N 1 49 32] # [2 N 2 49 32] # [2 N 4 49 32]
k1, v1 = kv1[0], kv1[1] #B head N C [N 1 49 32] [N 2 49 32] [N 4 49 32]attn1 = (q1 @ k1.transpose(-2, -1)) * self.scale #B head Nq Nkv [N 1 3136 49] [N 2 784 49] [N 4 196 49]
attn1 = attn1.softmax(dim=-1)
attn1 = self.attn_drop(attn1)
x1 = (attn1 @ v1).transpose(1, 2).reshape(B, N, C//2) # [N 3136 32] [N 784 64] [N 196 128]global_mask_value = torch.mean(attn1.detach().mean(1), dim=1) # B Nk  #max ?  mean ?  # [N 49]
global_mask_value = F.interpolate(global_mask_value.view(B,1,H//self.sr_ratio,W//self.sr_ratio),(H, W), mode='nearest')[:, 0]  # [N 56 56] [N 28 28] [N 14 14]# local
q2 = self.q2(x).reshape(B, N, self.num_heads // 2, C // self.num_heads).permute(0, 2, 1, 3) #B head N C # [N 1 3136 32] [N 2 784 32] [N 4 196 32]
kv2 = self.kv2(x_.reshape(B, C, -1).permute(0, 2, 1)).reshape(B, -1, 2, self.num_heads // 2,C // self.num_heads).permute(2, 0, 3, 1, 4)# [2 N 1 3136 32] [2 N 2 784 32] [2 N 4 196 32]
k2, v2 = kv2[0], kv2[1] # [N 1 3136 32] [N 2 784 32] [N 4 196 32]
q_window = 7
window_size= 7
q2, k2, v2 = window_partition(q2, q_window, H, W), window_partition(k2, window_size, H, W), \window_partition(v2, window_size, H, W) # [N*64,49,32] [N*32,49,32] [N*16 49 32]
attn2 = (q2 @ k2.transpose(-2, -1)) * self.scale # [N*64 49 49] [N*32 49 49] [N*16 49 49]
# (B*numheads*num_windows, window_size*window_size, window_size*window_size)
attn2 = attn2.softmax(dim=-1)
attn2 = self.attn_drop(attn2)x2 = (attn2 @ v2)  # B*numheads*num_windows, window_size*window_size, C   .transpose(1, 2).reshape(B, N, C) # [N*64 49 32] [N*32 49 32] [N*16 49 32]
x2 = window_reverse(x2, q_window, H, W, self.num_heads // 2) # [N 3136 32] [N 784 64] [N 196 128]local_mask_value = torch.mean(attn2.detach().view(B, self.num_heads//2, H//window_size*W//window_size, window_size*window_size, window_size*window_size).mean(1), dim=2) #[N 64 49]
local_mask_value = local_mask_value.view(B, H // window_size, W // window_size, window_size, window_size) # [N 8 8 7 7]
local_mask_value=local_mask_value.permute(0, 1, 3, 2, 4).contiguous().view(B, H, W) # [N 56 56] [N 28 28] [N 14 14]# mask B H W
x = torch.cat([x1, x2], dim=-1) # [N 3136 64] [N 784 128] [N 196 256]
x = self.proj(x+lepe) # linear(64,64) # linear(128,128) # linear(256,256)
x = self.proj_drop(x)
# cal mask
mask = local_mask_value+global_mask_value # [N 56 56] [N 28 28] [N 14 14]
mask_1 = mask.view(B, H * W) # [N 3136] [N 784] [N 196]
mask_2 = mask.permute(0, 2, 1).reshape(B, H * W) # [N 3136] [N 784] [N 196]
mask = [mask_1, mask_2]

1.4、Self-Guided Attention

尽管Self-Attention模型能够建模广泛的范围信息，但其高计算成本和内存消耗与序列长度的平方成正比，限制了其在各种计算机视觉任务中（如分割和检测）中对大尺寸特征图的适用性。最近的研究建议通过将一些 Token 合并成一个来减小序列长度。然而，这种聚合对待每个 Token 都是平等的，忽视了不同 Token 之间的固有重要性差异。这种聚合面临两个问题：

在显著区域，信息可能会丢失或与不相关的信息混合在一起
在次要区域或背景区域，大量 Token 对于简单的语义而言是多余的，但需要大量计算

受到这一观察的启发，我们提出了Self-Guided 注意力，它使用重要性作为引导来聚合 Token 。换句话说，在显著区域，保留更多的 Token 以获取细粒度的信息，而在次要区域，保留较少的 Token 以保持Self-Attention的全局视图并同时减少计算成本。

如图4所示，"self-guided"表示Transformer自己在训练期间确定了计算成本降低策略，而不是由人为引入的先验知识，如Swin中的窗口注意力，CSWin中的交叉形状注意力，PVT中的静态空间缩减。

输入特征图 $X\in R^{h\times w\times c}$ 首先投影为 Query (Q)，Key(K)和Value(V)。接下来，H个独立的Self-Attention Head并行计算Self-Attention。为了降低计算成本，同时保持Self-Attention之后特征图的大小不变，固定Q的长度，但使用重要性引导聚合模块(IAM)来聚合K和V的 Token 。

IAM的目标是在显著区域将更少的 Token 聚合为一个（即保留更多的 Token ），在背景区域将更多的 Token 聚合为一个（即保留更少的 Token ）。在公式（1）中，注意力图 $S\in R^{h\times w}$ 包含具有多粒度的区域重要性信息。

按升序对注意力图的Value进行排序，并将S均匀分成n个子区域 $s^{1}...s^{n}$ 。因此， $s^{n}$ 和 $s^{1}$ 分别是最重要的和次要的区域。同时，根据 $s^{1}...s^{n}$ 将 $X$ 中的所有 Token 分组为 $X^{1}...X^{n}$ 。在公式（1）中， $r$ 表示在不同重要性区域的聚合率，以便每个子区域都有一个聚合率，而子区域越重要，聚合率就越小。不同阶段的 $r$ 的具体值列在表1中。因此，IAM通过连接每个组的不同聚合率重新分配了每个组的分组输入特征 $X^{1}...X^{n}$ 的 Token 。

其中 $F(X,r)$ 是聚合函数，我们通过具有输入维度r和输出维度1的全连接层来实现它。 $\hat{X}^{i}$ 中的 Token 数量等于 $X^{i}$ 中的 Token 数量除以 $r_{i}$ 。对应这部分代码

q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # [N 2 3136 32] [N 4 784 32] [N 8 196 32]# mask [local_mask global_mask]  local_mask [value index]  value [B, H, W]
# use mask to fuse
mask_1, mask_2 = mask # [[N 3136],[N 3136]] [[N 784],[N 784]] [[N 196],[N 196]]
mask_sort1, mask_sort_index1 = torch.sort(mask_1, dim=1)
mask_sort2, mask_sort_index2 = torch.sort(mask_2, dim=1)
if self.sr_ratio == 8:token1, token2, token3 = H * W // (14 * 14), H * W // 56, H * W // 28  # [16 56 112]token1, token2, token3 = token1 // 4, token2 // 2, token3 // 4 # [4 28 28]
elif self.sr_ratio == 4:token1, token2, token3 = H * W // 49, H * W // 14, H * W // 7 # [16 56 112]token1, token2, token3 = token1 // 4, token2 // 2, token3 // 4 # [4 28 28]
elif self.sr_ratio == 2:token1, token2 = H * W // 2, H * W // 1 # [98 196]token1, token2 = token1 // 2, token2 // 2 # [49 98]
if self.sr_ratio==4 or self.sr_ratio==8:p1 = torch.gather(x, 1, mask_sort_index1[:, :H * W // 4].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C # [N 784 64] 根据mask中的index对x[:, :H * W // 4(784)]进行重新排序 [N 196 128]p2 = torch.gather(x, 1, mask_sort_index1[:, H * W // 4:H * W // 4 * 3].unsqueeze(-1).repeat(1, 1, C)) # [N 1568 64] [N 392 128]p3 = torch.gather(x, 1, mask_sort_index1[:, H * W // 4 * 3:].unsqueeze(-1).repeat(1, 1, C)) # [N 784 64] [N 196 128]seq1 = torch.cat([self.f1(p1.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1), # [N 64 4 196] # linear(196,1) # [N 128 4 49] # linear(49,1) 次要self.f2(p2.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1), # [N 64 28 56] # linear(56,1) # [N 128 28 14] # linear(14,1)self.f3(p3.permute(0, 2, 1).reshape(B, C, token3, -1)).squeeze(-1)], dim=-1).permute(0,2,1)  # B N C  # [N 64 28 28] # linear(28,1) # [N 128 28 7] # linear(7,1) 最重要# seq1 [N 60 64] # seq1 [N 60 128]x_ = x.view(B, H, W, C).permute(0, 2, 1, 3).reshape(B, H * W, C) # [N 3136 64] [N 784 128]p1_ = torch.gather(x_, 1, mask_sort_index2[:, :H * W // 4].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C  [N 784 64] [N 196 128]p2_ = torch.gather(x_, 1, mask_sort_index2[:, H * W // 4:H * W // 4 * 3].unsqueeze(-1).repeat(1, 1, C)) # [N 1568 64] [N 392 128]p3_ = torch.gather(x_, 1, mask_sort_index2[:, H * W // 4 * 3:].unsqueeze(-1).repeat(1, 1, C)) # [N 784 64] [N 196 128]seq2 = torch.cat([self.f1(p1_.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1),self.f2(p2_.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1),self.f3(p3_.permute(0, 2, 1).reshape(B, C, token3, -1)).squeeze(-1)], dim=-1).permute(0,2,1)  # B N C # seq2 [N 60 64] seq2 [N 60 128]
elif self.sr_ratio==2:p1 = torch.gather(x, 1, mask_sort_index1[:, :H * W // 2].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C [N 98 256]p2 = torch.gather(x, 1, mask_sort_index1[:, H * W // 2:].unsqueeze(-1).repeat(1, 1, C)) # [N 98 256]seq1 = torch.cat([self.f1(p1.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1), # [N 256 49 2] # linear(2,1)self.f2(p2.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1)], dim=-1).permute(0, 2, 1)  # B N C # [N 256 98 1] # linear(1,1)# seq1 [N 147 256]x_ = x.view(B, H, W, C).permute(0, 2, 1, 3).reshape(B, H * W, C)p1_ = torch.gather(x_, 1, mask_sort_index2[:, :H * W // 2].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C [N 98 256]p2_ = torch.gather(x_, 1, mask_sort_index2[:, H * W // 2:].unsqueeze(-1).repeat(1, 1, C)) # [N 98 256]seq2 = torch.cat([self.f1(p1_.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1),self.f2(p2_.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1)], dim=-1).permute(0, 2, 1)  # B N C# seq2 [N 147 256]
kv1 = self.kv1(seq1).reshape(B, -1, 2, self.num_heads // 2, C // self.num_heads).permute(2, 0, 3, 1, 4) # kv B heads N C # [2 N 1 60 32] # [2 N 2 60 32] # [2 N 4 147 32]
kv2 = self.kv2(seq2).reshape(B, -1, 2, self.num_heads // 2, C // self.num_heads).permute(2, 0, 3, 1, 4) # [2 N 1 60 32] # [2 N 2 60 32] # [2 N 4 147 32]
kv = torch.cat([kv1, kv2], dim=2) # [2 N 2 60 32] # [2 N 4 60 32] # [2 N 8 147 32]
k, v = kv[0], kv[1] # [N 2 60 32] # [N 4 60 32] # [N 8 147 32]
attn = (q @ k.transpose(-2, -1)) * self.scale # [N 2 3136 60] # [N 4 784 60] [N 8 196 147]
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)x = (attn @ v).transpose(1, 2).reshape(B, N, C) # [N 3136 64] # [N 784 128] [N 196 256]
x = self.proj(x+lepe)
x = self.proj_drop(x)
mask=None

将两部分的合起来看：

class Attention(nn.Module):def __init__(self, dim, mask, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1, linear=False):super().__init__()assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."self.dim = dimself.num_heads = num_headshead_dim = dim // num_headsself.scale = qk_scale or head_dim ** -0.5self.sr_ratio=sr_ratioif sr_ratio>1:if mask:self.q = nn.Linear(dim, dim, bias=qkv_bias)self.kv1 = nn.Linear(dim, dim, bias=qkv_bias)self.kv2 = nn.Linear(dim, dim, bias=qkv_bias)if self.sr_ratio==8:f1, f2, f3 = 14*14, 56, 28elif self.sr_ratio==4:f1, f2, f3 = 49, 14, 7elif self.sr_ratio==2:f1, f2, f3 = 2, 1, Noneself.f1 = nn.Linear(f1, 1)self.f2 = nn.Linear(f2, 1)if f3 is not None:self.f3 = nn.Linear(f3, 1)else:self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)self.norm = nn.LayerNorm(dim)self.act = nn.GELU()self.q1 = nn.Linear(dim, dim//2, bias=qkv_bias)self.kv1 = nn.Linear(dim, dim, bias=qkv_bias)self.q2 = nn.Linear(dim, dim // 2, bias=qkv_bias)self.kv2 = nn.Linear(dim, dim, bias=qkv_bias)else:self.q = nn.Linear(dim, dim, bias=qkv_bias)self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)self.lepe_linear = nn.Linear(dim, dim)self.lepe_conv = local_conv(dim)self.attn_drop = nn.Dropout(attn_drop)self.proj = nn.Linear(dim, dim)self.proj_drop = nn.Dropout(proj_drop)self.linear = linearself.apply(self._init_weights)def forward(self, x, H, W, mask):B, N, C = x.shapelepe = self.lepe_conv(self.lepe_linear(x).transpose(1, 2).view(B, C, H, W)).view(B, C, -1).transpose(-1, -2) # [N 3136 64] #[N 784 128] #[N 196 256] #[N 49 512]if self.sr_ratio > 1:if mask is None:# globalq1 = self.q1(x).reshape(B, N, self.num_heads//2, C // self.num_heads).permute(0, 2, 1, 3) # [N 1 3136 32] # [N 2 784 32] # [N 4 196 32]x_ = x.permute(0, 2, 1).reshape(B, C, H, W) # [N 64 56 56] # [N 128 28 28] # [N 256 14 14]x_1 = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1) # [N 49 64] sr --> conv2d(64, 64) # [N 49 128] sr --> conv2d(128, 128) # [N 49 256] sr --> conv2d(256, 256)x_1 = self.act(self.norm(x_1)) # [N 49 64] # [N 49 128] # [N 49 256]kv1 = self.kv1(x_1).reshape(B, -1, 2, self.num_heads//2, C // self.num_heads).permute(2, 0, 3, 1, 4) # [2 N 1 49 32] # [2 N 2 49 32] # [2 N 4 49 32]k1, v1 = kv1[0], kv1[1] #B head N C [N 1 49 32] [N 2 49 32] [N 4 49 32]attn1 = (q1 @ k1.transpose(-2, -1)) * self.scale #B head Nq Nkv [N 1 3136 49] [N 2 784 49] [N 4 196 49]attn1 = attn1.softmax(dim=-1)attn1 = self.attn_drop(attn1)x1 = (attn1 @ v1).transpose(1, 2).reshape(B, N, C//2) # [N 3136 32] [N 784 64] [N 196 128]global_mask_value = torch.mean(attn1.detach().mean(1), dim=1) # B Nk  #max ?  mean ?  # [N 49]global_mask_value = F.interpolate(global_mask_value.view(B,1,H//self.sr_ratio,W//self.sr_ratio),(H, W), mode='nearest')[:, 0]  # [N 56 56] [N 28 28] [N 14 14]# localq2 = self.q2(x).reshape(B, N, self.num_heads // 2, C // self.num_heads).permute(0, 2, 1, 3) #B head N C # [N 1 3136 32] [N 2 784 32] [N 4 196 32]kv2 = self.kv2(x_.reshape(B, C, -1).permute(0, 2, 1)).reshape(B, -1, 2, self.num_heads // 2,C // self.num_heads).permute(2, 0, 3, 1, 4)# [2 N 1 3136 32] [2 N 2 784 32] [2 N 4 196 32]k2, v2 = kv2[0], kv2[1] # [N 1 3136 32] [N 2 784 32] [N 4 196 32]q_window = 7window_size= 7q2, k2, v2 = window_partition(q2, q_window, H, W), window_partition(k2, window_size, H, W), \window_partition(v2, window_size, H, W) # [N*64,49,32] [N*32,49,32] [N*16 49 32]attn2 = (q2 @ k2.transpose(-2, -1)) * self.scale # [N*64 49 49] [N*32 49 49] [N*16 49 49]# (B*numheads*num_windows, window_size*window_size, window_size*window_size)attn2 = attn2.softmax(dim=-1)attn2 = self.attn_drop(attn2)x2 = (attn2 @ v2)  # B*numheads*num_windows, window_size*window_size, C   .transpose(1, 2).reshape(B, N, C) # [N*64 49 32] [N*32 49 32] [N*16 49 32]x2 = window_reverse(x2, q_window, H, W, self.num_heads // 2) # [N 3136 32] [N 784 64] [N 196 128]local_mask_value = torch.mean(attn2.detach().view(B, self.num_heads//2, H//window_size*W//window_size, window_size*window_size, window_size*window_size).mean(1), dim=2) #[N 64 49]local_mask_value = local_mask_value.view(B, H // window_size, W // window_size, window_size, window_size) # [N 8 8 7 7]local_mask_value=local_mask_value.permute(0, 1, 3, 2, 4).contiguous().view(B, H, W) # [N 56 56] [N 28 28] [N 14 14]# mask B H Wx = torch.cat([x1, x2], dim=-1) # [N 3136 64] [N 784 128] [N 196 256]x = self.proj(x+lepe) # linear(64,64) # linear(128,128) # linear(256,256)x = self.proj_drop(x)# cal maskmask = local_mask_value+global_mask_value # [N 56 56] [N 28 28] [N 14 14]mask_1 = mask.view(B, H * W) # [N 3136] [N 784] [N 196]mask_2 = mask.permute(0, 2, 1).reshape(B, H * W) # [N 3136] [N 784] [N 196]mask = [mask_1, mask_2]else:q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # [N 2 3136 32] [N 4 784 32] [N 8 196 32]# mask [local_mask global_mask]  local_mask [value index]  value [B, H, W]# use mask to fusemask_1, mask_2 = mask # [[N 3136],[N 3136]] [[N 784],[N 784]] [[N 196],[N 196]]mask_sort1, mask_sort_index1 = torch.sort(mask_1, dim=1)mask_sort2, mask_sort_index2 = torch.sort(mask_2, dim=1)if self.sr_ratio == 8:token1, token2, token3 = H * W // (14 * 14), H * W // 56, H * W // 28  # [16 56 112]token1, token2, token3 = token1 // 4, token2 // 2, token3 // 4 # [4 28 28]elif self.sr_ratio == 4:token1, token2, token3 = H * W // 49, H * W // 14, H * W // 7 # [16 56 112]token1, token2, token3 = token1 // 4, token2 // 2, token3 // 4 # [4 28 28]elif self.sr_ratio == 2:token1, token2 = H * W // 2, H * W // 1 # [98 196]token1, token2 = token1 // 2, token2 // 2 # [49 98]if self.sr_ratio==4 or self.sr_ratio==8:p1 = torch.gather(x, 1, mask_sort_index1[:, :H * W // 4].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C # [N 784 64] 根据mask中的index对x[:, :H * W // 4(784)]进行重新排序 [N 196 128]p2 = torch.gather(x, 1, mask_sort_index1[:, H * W // 4:H * W // 4 * 3].unsqueeze(-1).repeat(1, 1, C)) # [N 1568 64] [N 392 128]p3 = torch.gather(x, 1, mask_sort_index1[:, H * W // 4 * 3:].unsqueeze(-1).repeat(1, 1, C)) # [N 784 64] [N 196 128]seq1 = torch.cat([self.f1(p1.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1), # [N 64 4 196] # linear(196,1) # [N 128 4 49] # linear(49,1) 次要self.f2(p2.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1), # [N 64 28 56] # linear(56,1) # [N 128 28 14] # linear(14,1)self.f3(p3.permute(0, 2, 1).reshape(B, C, token3, -1)).squeeze(-1)], dim=-1).permute(0,2,1)  # B N C  # [N 64 28 28] # linear(28,1) # [N 128 28 7] # linear(7,1) 最重要# seq1 [N 60 64] # seq1 [N 60 128]x_ = x.view(B, H, W, C).permute(0, 2, 1, 3).reshape(B, H * W, C) # [N 3136 64] [N 784 128]p1_ = torch.gather(x_, 1, mask_sort_index2[:, :H * W // 4].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C  [N 784 64] [N 196 128]p2_ = torch.gather(x_, 1, mask_sort_index2[:, H * W // 4:H * W // 4 * 3].unsqueeze(-1).repeat(1, 1, C)) # [N 1568 64] [N 392 128]p3_ = torch.gather(x_, 1, mask_sort_index2[:, H * W // 4 * 3:].unsqueeze(-1).repeat(1, 1, C)) # [N 784 64] [N 196 128]seq2 = torch.cat([self.f1(p1_.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1),self.f2(p2_.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1),self.f3(p3_.permute(0, 2, 1).reshape(B, C, token3, -1)).squeeze(-1)], dim=-1).permute(0,2,1)  # B N C # seq2 [N 60 64] seq2 [N 60 128]elif self.sr_ratio==2:p1 = torch.gather(x, 1, mask_sort_index1[:, :H * W // 2].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C [N 98 256]p2 = torch.gather(x, 1, mask_sort_index1[:, H * W // 2:].unsqueeze(-1).repeat(1, 1, C)) # [N 98 256]seq1 = torch.cat([self.f1(p1.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1), # [N 256 49 2] # linear(2,1)self.f2(p2.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1)], dim=-1).permute(0, 2, 1)  # B N C # [N 256 98 1] # linear(1,1)# seq1 [N 147 256]x_ = x.view(B, H, W, C).permute(0, 2, 1, 3).reshape(B, H * W, C)p1_ = torch.gather(x_, 1, mask_sort_index2[:, :H * W // 2].unsqueeze(-1).repeat(1, 1, C))  # B, N//4, C [N 98 256]p2_ = torch.gather(x_, 1, mask_sort_index2[:, H * W // 2:].unsqueeze(-1).repeat(1, 1, C)) # [N 98 256]seq2 = torch.cat([self.f1(p1_.permute(0, 2, 1).reshape(B, C, token1, -1)).squeeze(-1),self.f2(p2_.permute(0, 2, 1).reshape(B, C, token2, -1)).squeeze(-1)], dim=-1).permute(0, 2, 1)  # B N C# seq2 [N 147 256]kv1 = self.kv1(seq1).reshape(B, -1, 2, self.num_heads // 2, C // self.num_heads).permute(2, 0, 3, 1, 4) # kv B heads N C # [2 N 1 60 32] # [2 N 2 60 32] # [2 N 4 147 32]kv2 = self.kv2(seq2).reshape(B, -1, 2, self.num_heads // 2, C // self.num_heads).permute(2, 0, 3, 1, 4) # [2 N 1 60 32] # [2 N 2 60 32] # [2 N 4 147 32]kv = torch.cat([kv1, kv2], dim=2) # [2 N 2 60 32] # [2 N 4 60 32] # [2 N 8 147 32]k, v = kv[0], kv[1] # [N 2 60 32] # [N 4 60 32] # [N 8 147 32]attn = (q @ k.transpose(-2, -1)) * self.scale # [N 2 3136 60] # [N 4 784 60] [N 8 196 147]attn = attn.softmax(dim=-1)attn = self.attn_drop(attn)x = (attn @ v).transpose(1, 2).reshape(B, N, C) # [N 3136 64] # [N 784 128] [N 196 256]x = self.proj(x+lepe)x = self.proj_drop(x)mask=Noneelse:q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # [N 16 49 32]kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) #[2 N 16 49 32]k, v = kv[0], kv[1] # [N 16 49 32]attn = (q @ k.transpose(-2, -1)) * self.scale # [N 16 49 49]attn = attn.softmax(dim=-1)attn = self.attn_drop(attn)x = (attn @ v).transpose(1, 2).reshape(B, N, C) # [N 49 512]x = self.proj(x+lepe) # linear(512,512)x = self.proj_drop(x)mask=Nonereturn x, mask

其中分窗和swin transformer是一样的。

在分类上的表现：