当前位置：首页 > news >正文

深度学习·经典模型·SwinTransformer

news 来源：原创 2025/4/30 15:36:10

SwinTransformer

主要创新点：移动窗口，基于窗口的注意力计算

Patch Embedding

下采样打包为Pacth：可以直接使用Conv2d
也可以先打包后使用embedding映射。

Patch Merging

类似池化的操作,压缩图片大小，同时通道数增多，获得更多的语义信息。
实现：获得相邻的Patch,然后在通道维度上concat，维度变为 $4 C$ ，最后经过线性层投射回 $2 C$ 。

例子

1,2,5,6是相邻的2x2的Patch

[	[1,2,3,4][5,6,7,8][9,10,11,12][13,14,15,16]]

通过切片获得对应位置的元素
注意我们从通道维度上拼接，所以不能按照传统的上下拼接的思路理解
这段代码的效果是：编号1,2,5,6的特征向量拼接，(相邻元素就好像叠加在一起)

        x0 = x[:, 0::2, 0::2, :]  # [B, H/2, W/2, C]x1 = x[:, 1::2, 0::2, :]  # [B, H/2, W/2, C]x2 = x[:, 0::2, 1::2, :]  # [B, H/2, W/2, C]x3 = x[:, 1::2, 1::2, :]  # [B, H/2, W/2, C]x = torch.cat([x0, x1, x2, x3], -1)  # [B, H/2, W/2, 4*C]

窗口注意力机制

本文使用了窗口注意力机制，计算复杂度是随着图像扩大线性增长的
这点就好比 $8 * 8 = 64 >>> 4 * (2 * 2) = 16$ )，小窗口 $2 * 2$ 比大窗口 $8 * 8$ 明显复杂度低很多。

实现讲解

输入： $B*N_w,M_h*M_w,C)$ ： $N_w$ 是窗口数量，可以参加Embedding层 $N_w=\frac{H*W}{M_h*M_w}$
输入的理解：将窗口数量理解为一种批次， $M_h*M_w$ 作为序列的长度，reshape为指定维度： $B*N_w,M_h*N_w,C)$
快速计算KQV，直接使用线性层映射为 $B*N_w,M_h*N_w,3C)$ ，然后拆分最后一个维度 $3 C$ ，变成各自 $3，B*N_w,M_h*N_w,C)$ 的QKV大小，为分离QKV作准备。
多头注意力机制：每一个KQV维度 $3,B*N_w,M_h*N_w,C)$ ，转换为 $3，B*N_w,N_{head},M_h*N_w,dim_{head})$ , $N_{head}$ 不会参与计算，只需要最后两个维度进行KQV的矩阵乘法即可获得最终的多头注意力输出!
然后就是Masked掩码操作：这里使用的是加性掩码，掩码的生成方式见下。
输出维度不变： $B*N_w,M_h*M_w,C)$

    def forward(self, x, mask: Optional[torch.Tensor] = None):"""Args:x: input features with shape of (num_windows*B, Mh*Mw, C)mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None"""# [batch_size*num_windows, Mh*Mw, total_embed_dim]B_, N, C = x.shape# qkv(): -> [batch_size*num_windows, Mh*Mw, 3 * total_embed_dim]# reshape: -> [batch_size*num_windows, Mh*Mw, 3, num_heads, embed_dim_per_head]# permute: -> [3, batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)# [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)# transpose: -> [batch_size*num_windows, num_heads, embed_dim_per_head, Mh*Mw]# @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, Mh*Mw]q = q * self.scaleattn = (q @ k.transpose(-2, -1))# relative_position_bias_table.view: [Mh*Mw*Mh*Mw,nH] -> [Mh*Mw,Mh*Mw,nH]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)relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # [nH, Mh*Mw, Mh*Mw]attn = attn + relative_position_bias.unsqueeze(0)if mask is not None:# mask: [nW, Mh*Mw, Mh*Mw]nW = mask.shape[0]  # num_windows# attn.view: [batch_size, num_windows, num_heads, Mh*Mw, Mh*Mw]# mask.unsqueeze: [1, nW, 1, Mh*Mw, Mh*Mw]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)else:attn = self.softmax(attn)attn = self.attn_drop(attn)# @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]# transpose: -> [batch_size*num_windows, Mh*Mw, num_heads, embed_dim_per_head]# reshape: -> [batch_size*num_windows, Mh*Mw, total_embed_dim]x = (attn @ v).transpose(1, 2).reshape(B_, N, C)x = self.proj(x)x = self.proj_drop(x)return x

滑动窗口的实现

本文的核心操作：实现起来不难
实现代码：注意图像整体往右下，roll这个函数是相当于移动窗口的，所以是往左上移动窗口
输入和输出是以图片的格式： $(B, H * W, C)$

        if self.shift_size > 0:shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))

移动窗口注意力

先调用移动窗口：对图像进行移动处理。
使用被移动后的图像进行窗口注意力计算,输出维度 $B*N_w,M_h*M_w,C)$ ：
还原为图像 $(B, H, W, C)$
以反方向移动图像：x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
reshape为： $(B, H * W, C)$ ，丢入MLP中处理，放大 $4 C$ ，然后还原为 $C$ 。

在这里插入图片描述

MASK的实现

建议直接抄以下代码：

    def create_mask(self, x, H, W):# calculate attention mask for SW-MSA# 保证Hp和Wp是window_size的整数倍Hp = int(np.ceil(H / self.window_size)) * self.window_sizeWp = int(np.ceil(W / self.window_size)) * self.window_size# 拥有和feature map一样的通道排列顺序，方便后续window_partitionimg_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # [1, Hp, Wp, 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 = 0for h in h_slices:for w in w_slices:img_mask[:, h, w, :] = cntcnt += 1mask_windows = window_partition(img_mask, self.window_size)  # [nW, Mh, Mw, 1]mask_windows = mask_windows.view(-1, self.window_size * self.window_size)  # [nW, Mh*Mw]attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)  # [nW, 1, Mh*Mw] - [nW, Mh*Mw, 1]# [nW, Mh*Mw, Mh*Mw]attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))return attn_mask