Swin Transformer代码阅读注释

在这里插入图片描述

Swin Transformer代码阅读注释

前言

上一篇博文以论文中的内容介绍了Swin Transformer的网络结构和一些细节。本篇博文将从官方代码中的

swin_transformer.py

去详细介绍Swin Transformer结构，并补充代码中才有而论文中没有的细节。

|| 如果对 Swin Transformer 不了解建议先看论文介绍再看源码 ||
Swin Transformer介绍博客：论文阅读笔记：Swin Transformer

Swin Transformer

代码中实现的网络结构与论文中的结构如如下：
在这里插入图片描述
如上图所示，代码中使用

PatchEmbed

来实现

Patch Partition

Linear Embedding

,使用

BasicLayer

来实现

Swin Transformer Block

PatchMerging

，对于最后一个

BasicLayer

不使用

PatchMerging

来降采样。

Swin-T 的配置如下：
Swin-T参数配置表
在这里插入图片描述
网络结构介绍可看Swin Transformer介绍博客：论文阅读笔记：Swin Transformer
整体结构代码和注释如下（代码大部分和和 Vision Transformer 是一样）：

classSwinTransformer(nn.Module):def__init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
                 embed_dim=96, depths=[2,2,6,2], num_heads=[3,6,12,24],
                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
                 use_checkpoint=False,**kwargs):'''
        img_size(int | tuple(int)): 输入图像尺寸. 默认: 224
        patch_size (int | tuple(int)): Patch尺寸. 默认: Swin-T参数配置表中 stage1中的96
        in_chans (int): 输入图像通道. 默认: 3
        num_classes (int): 分类数. 默认: 1000
        embed_dim (int): Patch embedding的输出通道. 默认: 96 (Swin-T参数配置表中 stage 1 中的 96-d)
        depths (tuple(int)):Swin Transformer Block 的个数. 默认:[2, 2, 6, 2] (Swin-T参数配置表中的[×2, ×2, ×6, ×2])
        num_heads (tuple(int)): 不同层 MSA 计算中的 head 数. 默认：[3, 6, 12, 24] (Swin-T参数配置表中的[head 3，head 6，head 12，head 24])
        window_size (int): W-MSA 和 SW-MSA 的 Window 尺寸. 默认: 7 (Swin-T参数配置表中的 “win.sz. 7×7”)
        mlp_ratio (float): 通过MLP的输出通道倍数. 默认: 4 (Swin-T参数配置表中的“dim 96”,“dim 192”,“dim 384”,“dim 768”可以看出)
        qkv_bias (bool): 使用 Linear 将输入映射到 qkv 时，Linear是否使用偏置. 默认: True
        qk_scale (float):qk缩放比例，如果是 None 则使用根号 dim_k 分之一. 默认: None
        drop_rate (float): dropout概率. 默认: 0
        attn_drop_rate (float): attention 中的 dropout 概率. 默认: 0
        drop_path_rate (float): attention 中的 droppath 概率. 默认: 0.1
        norm_layer (nn.Module): 归一化方式. 默认: nn.LayerNorm.
        ape (bool): 是否在 patch embedding 后使用绝对位置编码. 默认: False
        patch_norm (bool): 是否在 patch embedding 后使用归一化. 默认: True
        use_checkpoint (bool): 是否 checkpointing 节省内存. 默认: False
        '''super().__init__()
        
        self.num_classes = num_classes
        self.num_layers =len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        '''
        经过4个stage后的通道数（从96->768 即：96*2^(4-1)=768）
        '''
        self.num_features =int(embed_dim *2**(self.num_layers -1))
        self.mlp_ratio = mlp_ratio
        '''
        将图片划分成没有重叠的多个patch
        PatchEmbed代码在下文中介绍
        patches_resolution = [img_size[0]//patch_size[0], img_size[1]//patch_size[1]] = [56,56]
        num_patches = patches_resolution[0] * patches_resolution[1] = 3136
        '''
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm elseNone)
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution
        '''
        如果使用绝对位置编码则构建可学习的绝对位置编码参数：
        self.absolute_pos_embed : [1,3136,96]
        默认不使用
        '''if self.ape:
            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
            trunc_normal_(self.absolute_pos_embed, std=.02)'''
        pos_drop 以 drop_rate 概率进行 Dropout
        '''
        self.pos_drop = nn.Dropout(p=drop_rate)'''
        构建首项为0，长度为depths(2+2+6+2=12)的等差数列，且最后一项小于drop_path_rate
        也就是说 传入 BasicLayer 的 droppath 概率是递增的。
        代码这里是让 drop_path_ratio 默认等于0.1
        最后利用参数构建 depth(12) 层 BasicLayer 层 
        BasicLayer 的代码在下文中介绍
        '''
        dpr =[x.item()for x in torch.linspace(0, drop_path_rate,sum(depths))]# stochastic depth decay rule# build layers
        self.layers = nn.ModuleList()for i_layer inrange(self.num_layers):
            layer = BasicLayer(dim=int(embed_dim *2** i_layer),#每个Basiclayer模块后通道数都翻倍'''
                               每个Basiclayer都进行了降采样
                               所以input_resolution每一次都要除以2
                               '''
                               input_resolution=(patches_resolution[0]//(2** i_layer), 
                                                 patches_resolution[1]//(2** i_layer)), 
                               depth=depths[i_layer],
                               num_heads=num_heads[i_layer],
                               window_size=window_size,
                               mlp_ratio=self.mlp_ratio,
                               qkv_bias=qkv_bias, qk_scale=qk_scale,
                               drop=drop_rate, attn_drop=attn_drop_rate,
                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer +1])],
                               norm_layer=norm_layer,'''
                               如果i_layer 是最后一层则不使用PatchMerging 来降采样
                               '''
                               downsample=PatchMerging if(i_layer < self.num_layers -1)elseNone,
                               use_checkpoint=use_checkpoint)
            self.layers.append(layer)'''
        进行归一化和平均池化
        最后用一个Linear做预测head
        '''
        self.norm = norm_layer(self.num_features)
        self.avgpool = nn.AdaptiveAvgPool1d(1)
        self.head = nn.Linear(self.num_features, num_classes)if num_classes >0else nn.Identity()
        self.apply(self._init_weights)'''
    初始化权重
    '''def_init_weights(self, m):ifisinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)ifisinstance(m, nn.Linear)and m.bias isnotNone:
                nn.init.constant_(m.bias,0)elifisinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias,0)
            nn.init.constant_(m.weight,1.0)defforward_features(self, x):'''
        如图所示先进行patch embedding
        如果使用绝对位置偏置就加上绝对位置编码
        '''
        x = self.patch_embed(x)if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)'''
        循环执行Blocks
        '''for layer in self.layers:
            x = layer(x)'''
        归一化并平均池化
        '''
        x = self.norm(x)# B L C
        x = self.avgpool(x.transpose(1,2))# B C 1
        x = torch.flatten(x,1)return x
    defforward(self, x):
        x = self.forward_features(x)'''
        对swin transformer的特征提取进行预测
        '''
        x = self.head(x)return x

1 PatchEmbed

Swin Transformer中的

PatchEmbed

模块和 VIT 中的 Linear Projection of Flattened Patches:

PatchEmbed

模块差不多，可查看博文Vision Transformer（Pytorch版）代码阅读注释查看，其主要思想是通过感受野大小等于步距大小的卷积来实现，与 VIT 不同的是其使用了

nn.LayerNorm

。

PatchEmbed

代码和注释如下：

classPatchEmbed(nn.Module):""" Image to Patch Embedding
    Args:
        img_size (int): 图像尺寸.  默认: 224.
        patch_size (int): token尺寸. 默认: 4.
        in_chans (int): 图像通道. 默认: 3.
        embed_dim (int): patch embed通道. 默认: 96.
        norm_layer (nn.Module, optional): 归一化. Default: None
    """def__init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):super().__init__()'''
        self.image_size = (224,224)
        self.patch_size = (4,4)
        self.patches_resolution = [56,56]
        self.num_patches = 56*56=3136
        '''
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        patches_resolution =[img_size[0]// patch_size[0], img_size[1]// patch_size[1]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0]* patches_resolution[1]'''
        self.in_chans = 3
        self.embed_dim = 96
        '''
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        
        '''
        self.proj = nn.Conv2d(3,96,(4,4),4)
        self.norm = nn.LayerNorm(96)
        '''
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)if norm_layer isnotNone:
            self.norm = norm_layer(embed_dim)else:
            self.norm =Nonedefforward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraintsassert H == self.img_size[0]and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."'''
        self.proj(x):[B,3,224,224]->[B,96,56,56]
        flatten(2):[B,96,56,56]->[B,96,56*56]=[B,96,3136]
        transpose(1, 2):[B,96,3136]->[B,3136,96]
        self.norm(x):[B,3136,96]->[B,3136,96]
        '''    
        x = self.proj(x).flatten(2).transpose(1,2)# B Ph*Pw Cif self.norm isnotNone:
            x = self.norm(x)return x

BasicLayer

代码中使用

BasicLayer

来实现论文中的

Swin Transformer Block

PatchMerging

，对于最后一个

BasicLayer

不使用

PatchMerging

来降采样。

BasicLayer

的代码和注释如下：

classBasicLayer(nn.Module):""" A basic Swin Transformer layer for one stage.
    Args:
        dim (int): 输入特征图的通道数.
        input_resolution (tuple[int]): 输入特征图的分辨率大小.
        depth (int): SwinTransformerBlock的个数.
        num_heads (int): Muti-Head Self-Attention 中的head个数.
        window_size (int): window 大小.
        mlp_ratio (float): patch embedding通过MLP的通道倍数.
        qkv_bias (bool): 使用 Linear 将输入映射到 qkv 时，Linear是否使用偏置. 默认: True
        qk_scale (float):qk缩放比例，如果是 None 则使用根号 dim_k 分之一. 默认: None
        drop (float, optional): dropout概率. 默认: 0
        attn_drop (float, optional): attention 中的 dropout 概率. 默认: 0
        drop_path (float | tuple[float], optional): attention 中的 droppath 概率. 默认: 0.1
        norm_layer (nn.Module): 归一化方式. 默认: nn.LayerNorm.
        downsample (nn.Module | None, optional): 降采样层. 默认: None 代码使用PatchMerging
        use_checkpoint (bool): 是否 checkpointing 节省内存. 默认: False
    """def__init__(self, dim, input_resolution, depth, num_heads, window_size,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint
        '''
        构建SwinTransformerBlock
        SwinTransformerBlock代码在下文介绍
        '''
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
                                 num_heads=num_heads, window_size=window_size,'''
                                 如果i是偶数，则表示是W-MSA，shift_size =0
                                 如果i是奇数，则表示是SW-MSA，shift_size = window_size // 2
                                 '''
                                 shift_size=0if(i %2==0)else window_size //2,
                                 mlp_ratio=mlp_ratio,
                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
                                 drop=drop, attn_drop=attn_drop,
                                 drop_path=drop_path[i]ifisinstance(drop_path,list)else drop_path,
                                 norm_layer=norm_layer)for i inrange(depth)])'''
        使用PatchMerging进行降采样
        PatchMerging代码在下文介绍
        '''if downsample isnotNone:
            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)else:
            self.downsample =Nonedefforward(self, x):for blk in self.blocks:if self.use_checkpoint:
                x = checkpoint.checkpoint(blk, x)else:
                x = blk(x)if self.downsample isnotNone:
            x = self.downsample(x)return x

SwinTransformerBlock

SwinTransformerBlock

的结构如下：
在这里插入图片描述

Mlp

此处和 VIT 中

MLP

一模一样，可查看Vision Transformer（Pytorch版）代码阅读注释。代码也很简单，就不再做任何赘述了。

Mlp

的代码如下：

classMlp(nn.Module):def__init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)defforward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)return x

window_partition

W-MSA

和

SW-MSA

首先需要将特征图拆分成多个windows。
在这里插入图片描述

window_partition

的代码和注释如下：

defwindow_partition(x, window_size):"""
    Args:
        x: (B, H, W, C)
        window_size (int): window size
    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    '''
    [B, H, W, C] -> [BHW//(window_size*window_size), window_size, window_size, C]
    '''
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0,1,3,2,4,5).contiguous().view(-1, window_size, window_size, C)return windows

window_reverse

在对每一个windows进行

WSA

计算以后需要将其还原成正常的特征图传入下一模块中。其实就是

window_partition

的逆过程。

window_reverse

的代码和注释如下：

defwindow_reverse(windows, window_size, H, W):"""
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image
    Returns:
        x: (B, H, W, C)
    """
    B =int(windows.shape[0]/(H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size,-1)
    x = x.permute(0,1,3,2,4,5).contiguous().view(B, H, W,-1)return x

WindowAttention

WindowAttention就是在 Vision Transformer 模块的Attention基础上加入了相对位置偏移

relative_position_bias_table

（即论文中提出的 Relative Position Bias）来提升精度：
在这里插入图片描述
生成相对位置偏置的过程（以一个head为例，假设

window_h = 2,window_w = 2

,相关代码已在图中标出）：
1.随机生成相对位置偏置表

relative_position_bias_table

：
在这里插入图片描述

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

2.首先windows内部每个像素都有自己的位置编码，其绝对位置编码的坐标

abs_coords

如果以左上角为原点，则如下图：
在这里插入图片描述

coords_h = torch.arange(self.window_size[0])
coords_w = torch.arange(self.window_size[1])
coords = torch.stack(torch.meshgrid([coords_h, coords_w]))# 2, Wh, Ww
coords_flatten = torch.flatten(coords,1)# 2, Wh*Ww

3.为了获得每个绝对坐标相对于其他坐标的相对位置，则需要用每个绝对坐标减去其他绝对坐标，即:

用 ( 0 , 0 ) (0,0) (0,0) 分别减去 ( 0 , 1 ) (0,1) (0,1)， ( 1 , 0 ) (1,0) (1,0)， ( 1 , 1 ) (1,1) (1,1)
用 ( 0 , 1 ) (0,1) (0,1) 分别减去 ( 0 , 0 ) (0,0) (0,0)， ( 1 , 0 ) (1,0) (1,0)， ( 1 , 1 ) (1,1) (1,1)
用 ( 1 , 0 ) (1,0) (1,0) 分别减去 ( 0 , 0 ) (0,0) (0,0)， ( 0 , 1 ) (0,1) (0,1)， ( 1 , 1 ) (1,1) (1,1)
用 ( 1 , 1 ) (1,1) (1,1) 分别减去 ( 0 , 0 ) (0,0) (0,0)， ( 0 , 1 ) (0,1) (0,1)， ( 1 , 0 ) (1,0) (1,0)

以此得到相对坐标

relative_coords

。这样每一个像素相对于其他像素的相对位置关系就得到了。
代码中实现过程如下：
在这里插入图片描述

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

4.将二维的坐标转换成一维的索引

relative_position_index

。

先将每一行加上 w i n d o w − h − 1 window_-h-1 window−h−1,每一列加上 w i n d o w − w − 1 window_-w-1 window−w−1
然后将每一行乘上 2 ∗ w i n d o w − w − 1 2*window_-w-1 2∗window−w−1
最后将行列坐标相加，得到relative_position_index

在这里插入图片描述

relative_coords[:,:,0]+= self.window_size[0]-1# shift to start from 0
relative_coords[:,:,1]+= self.window_size[1]-1
relative_coords[:,:,0]*=2* self.window_size[1]-1
relative_position_index = relative_coords.sum(-1)# Wh*Ww, Wh*Ww

这里作用的用意是为了让不同的相对关系只对应一种数字。如果一开始直接将行列坐标相加，会发现

    (
   
   
    0
   
   
    ,
   
   
    1
   
   
    )
   
  
  
   (0,1)
  
 
(0,1) 和 
 
  
   
    (
   
   
    1
   
   
    ,
   
   
    0
   
   
    )
   
  
  
   (1,0)
  
 
(1,0) 这两种关系无法分辨（相加都是1）。而如果不乘上 
 
  
   
    2
   
   
    ∗
   
   
    w
   
   
    i
   
   
    n
   
   
    d
   
   
    o
   
   
    
     w
    
    
     −
    
   
   
    w
   
   
    −
   
   
    1
   
  
  
   2*window_-w-1
  
 
2∗window−w−1，会发现
 
  
   
    (
   
   
    1
   
   
    ,
   
   
    1
   
   
    )
   
  
  
   (1,1)
  
 
(1,1) 和 
 
  
   
    (
   
   
    2
   
   
    ,
   
   
    0
   
   
    )
   
  
  
   (2,0)
  
 
(2,0) 这两种关系无法分辨（相加都是2）。所以经过几次变换以后再转换成1维的坐标就可以得到唯一的相对位置关系，如下图：

在这里插入图片描述

0代表右下的相对位置关系
1代表正下的相对位置关系
2代表坐下的相对位置关系
3代表正右的相对位置关系
4代表和自己的相对位置关系
5代表正左的相对位置关系
6代表右上的相对位置关系
7代表正上的相对位置关系
8代表左上的相对位置关系

5.通过相对位置索引

relative_position_index

在相对位置偏置表

relative_position_bias_table

中取找对应的相对位置偏置

relative_position_bias

。
在这里插入图片描述

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

WindowAttention

的代码和注释如下（大部分代码在Vision Transformer（Pytorch版）代码阅读注释的Multi-Head Attention部分已经介绍，主要还是相对位置偏置部分不同）：

classWindowAttention(nn.Module):""" Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    Args:
        dim (int): 输入特征图的通道
        window_size (tuple[int]): window的尺寸.
        num_heads (int): muti-head self-attention的head个数.
        qkv_bias (bool, optional):  是否使用 qkv 偏置（即使用 Linear 将输入映射到 qkv 时，Linear是否使用 bias ）
        qk_scale (float | None, optional): qk缩放比例，默认使用根号 dim_k 分之一
        attn_drop (float, optional): attention 中的 dropout 概率
        proj_drop (float, optional): linear 中的 dropout 概率
    """def__init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):'''
        VIT源码阅读中已经讲过了
        '''super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim **-0.5'''
        上述生成相对位置偏置的过程对应的代码：
        1.生成相对位置偏置表
        relative_position_bias_table.shape = [2*window_h-1 * 2*window_w-1, num_heads]
        '''
        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'''
        2.生成绝对位置坐标
        coords.shape = [2, window_h, window_w]
        '''
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))# 2, Wh, Ww'''
        3.生成相对位置坐标
        relative_coords.shape  = [window_h*window_w, window_h*window_w, 2]
        '''
        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'''
        4.将二维相对位置坐标转换成一维相对位置索引
        relative_position_index  = [window_h*window_w, window_h*window_w]
        '''
        relative_coords[:,:,0]+= self.window_size[0]-1# shift to start from 0
        relative_coords[:,:,1]+= self.window_size[1]-1
        relative_coords[:,:,0]*=2* self.window_size[1]-1
        relative_position_index = relative_coords.sum(-1)# Wh*Ww, Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)'''
        VIT源码阅读中已经讲过了,是一模一样的
        '''
        self.qkv = nn.Linear(dim, dim *3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)defforward(self, x, mask=None):"""
        Args:
            x: shape为(num_windows*B, N, C)的特征图
            mask: 为了使SW-MSA中不相邻的子窗口之间的不进行qk匹配，论文阅读笔记中已介绍，
            mask的生成和使用原理在后面介绍SwinTransformerBlock的代码时会介绍
        """
        B_, N, C = x.shape
        '''
        将VIY和attention时已经讲过了，不再介绍
        '''
        qkv = self.qkv(x).reshape(B_, N,3, self.num_heads, C // self.num_heads).permute(2,0,3,1,4)
        q, k, v = qkv[0], qkv[1], qkv[2]# make torchscript happy (cannot use tensor as tuple)
        q = q * self.scale
        attn =(q @ k.transpose(-2,-1))'''
        5.根据相对位置索引在相对位置偏置表中找到对应的相对位置偏置并加上
        '''
        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)'''
        利用mask掩膜计算将不相邻的子窗口之间使用softmax抑制来去除qk匹配
        '''if mask isnotNone:
            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)else:
            attn = self.softmax(attn)
        attn = self.attn_drop(attn)
        
        x =(attn @ v).transpose(1,2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)return x

Block

在论文阅读笔记：Swin Transformer中讲过Efficient batch computation for shifted configuration：
在这里插入图片描述
其将A、B、C三个框中的 window 移动到四个 4 × 4 红色框的对应位置，使其凑成四个 4 × 4 的window。由于有几个 window 是由不相邻的子窗口组成，需要通过Masked MSA 掩膜计算来限制每个 window 中的不同子窗口的 MSA。

掩膜计算相关代码如下：

'''
self.input_resolution为当前特征图的大小
img_mask 用于记录同特征图上每个像素的掩膜权重
h_slices w_slices 分别为宽高切片
self.shift_size 只有为SW-MSA计算时为self.window_size//2，W-MSA计算时为0
'''
H, W = self.input_resolution
img_mask = torch.zeros((1, H, W,1))# 1 H W 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))'''
将不同区域按0到8进行标记
''' 
cnt =0for h in h_slices:for w in w_slices:
        img_mask[:, h, w,:]= cnt
        cnt +=1
        mask_windows = window_partition(img_mask, self.window_size)# nW, window_size, window_size, 1
        mask_windows = mask_window s.view(-1, self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1)- mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask !=0,float(-100.0)).masked_fill(attn_mask ==0,float(0.0))

过程如下（假设

self.input_resolution=(8,8)

，

self.window_size=4

，

self.shift_size=self.window_size//2=2

）：
1.按特征图中的 window 通过切片方式生成划分子窗口的

img_mask

：
在这里插入图片描述
2.对

img_mask

中不同区域按0到8进行标记：
在这里插入图片描述
3.将window拆分并展开：

4.每个window中分别用每个像素的标记值互减，计算像素之间是否在同一个子window中得到

attn_mask(attn_mask.shape=[4,16,16])

：

attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)

举例说明：
第一个window标记向量为:

[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]

，依次减去每一个标记值，进行了16×16次减法，得到16×16的矩阵

[[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],...,[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]

第二个window标记向量为:

[1,1,2,2,1,1,2,2,1,1,2,2,1,1,2,2]

，依次减去每一个标记值，进行了16×16次减法，得到16×16的矩阵

[[0,0,1,1,0,0,1,1,0,0,1,1,0,0,1,1],...,[-1,-1,0,0,-1,-1,0,0,-1,-1,0,0,-1,-1,0,0]]

第三个window标记向量为:

[3,3,3,3,3,3,3,3,6,6,6,6,6,6,6,6]

，依次减去每一个标记值，进行了16×16次减法，得到16×16的矩阵

[[0,0,0,0,0,0,0,0,3,3,3,3,3,3,3,3],...,[0,0,0,0,0,0,0,0,-3,-3,-3,-3,-3,-3,-3,-3]]

第四个window标记向量为:

[4,4,5,5,4,4,5,5,7,7,8,8,7,7,8,8]

，依次减去每一个标记值，进行了16×16次减法，得到16×16的矩阵

[[0,0,1,1,0,0,1,1,3,3,4.4,3,3,4,4],...,[-4,-4,-3,-3,-4,-4,-3,-3,-1,-1,0,0,-1,-1,0,0]]

attn_mask

中某个window第

    i
   
  
  
   i
  
 
i 行第 
 
  
   
    j
   
  
  
   j
  
 
j 列如果为 0，则表示第 
 
  
   
    i
   
  
  
   i
  
 
i 个像素和第 
 
  
   
    j
   
  
  
   j
  
 
j 个像素属于同一个子窗口。

5.给

attn_mask

赋值（如果为0，则等于0，不为0，则等于-100）：

attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

结合

WindowAttention

的代码：

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)

其实现原理如下（以第0个像素为例）：

第0个像素与其他像素的再完成

    Q
   
   
    
     K
    
    
     T
    
   
   
    /
   
   
    
     d
    
   
   
    +
   
   
    B
   
  
  
   QK^T/\sqrt{d}+B
  
 
QKT/d+B的计算后，得到 
 
  
   
    
     α
    
    
     
      0
     
     
      ,
     
     
      0
     
    
   
  
  
   \alpha_{0,0}
  
 
α0,0 ,
 
  
   
    
     α
    
    
     
      0
     
     
      ,
     
     
      1
     
    
   
  
  
   \alpha_{0,1}
  
 
α0,1…
 
  
   
    
     α
    
    
     
      0
     
     
      ,
     
     
      15
     
    
   
  
  
   \alpha_{0,15}
  
 
α0,15,再进行掩膜计算

attn = attn + mask

，这样与像素0不在同一子窗口的像素结果就会减去100，最后通过softmax使其等于0，这样就抑制了不同窗口间的qk匹配。
再进行掩膜计算，
在这里插入图片描述
那么代码是如何实现子窗口之间的移动的呢？
窗口移动代码如下：

'''
X.shape = [B,H,W,C]
在 H 所在维度上滚动-self.shift_size个像素，
在 W 所在维度滚动-self.shift_size个像素
shift_size = window_size//2
注：torch.roll正方向是往下
'''if self.shift_size >0:
    shifted_x = torch.roll(x, shifts=(-self.shift_size,-self.shift_size), dims=(1,2))else:
    shifted_x = x

其实现过程如下(假设

window_size = 4

shift_size = window_size //2 = 2

)：
在这里插入图片描述
接着对图像进行拆分即可：

SwinTransformerBlock

的代码和注释如下：

classSwinTransformerBlock(nn.Module):""" Swin Transformer Block.
    Args:
        dim (int): 特征图的维度.
        input_resolution (tuple[int]): 输入特征图的分辨率.
        num_heads (int): muti-head self-attention的headg个数.
        window_size (int): 窗口尺寸.
        shift_size (int): SW-MSA的偏移尺寸，为window_size // 2.
        mlp_ratio (float): patch embedding通过MLP的通道倍数.
        qkv_bias (bool): 使用 Linear 将输入映射到 qkv 时，Linear是否使用偏置. 默认: True
        qk_scale (float):qk缩放比例，如果是 None 则使用根号 dim_k 分之一. 默认: None
        drop (float, optional): dropout概率. 默认: 0
        attn_drop (float, optional): attention 中的 dropout 概率. 默认: 0
        drop_path (float | tuple[float], optional): attention 中的 droppath 概率. 默认: 0.1
        act_layer (nn.Module, optional): 激活函数. 默认: nn.GELU
        norm_layer (nn.Module, optional): 归一化层.  默认: nn.LayerNorm
        
    """def__init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        '''
        如果窗口尺寸已经大于分辨率，则让窗口尺寸等于分辨率H、W 中较小的那一个
        '''ifmin(self.input_resolution)<= self.window_size:# if window size is larger than input resolution, we don't partition windows
            self.shift_size =0
            self.window_size =min(self.input_resolution)assert0<= self.shift_size < self.window_size,"shift_size must in 0-window_size"
        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.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)'''
        如果偏移量大于0，则表示使用SW-MSA
        计算attn_mask,原理如上述过程。
        '''if self.shift_size >0:# calculate attention mask for SW-MSA
            H, W = self.input_resolution
            img_mask = torch.zeros((1, H, W,1))# 1 H W 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,:]= cnt
                    cnt +=1
            mask_windows = window_partition(img_mask, self.window_size)# nW, window_size, window_size, 1
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1)- mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask !=0,float(-100.0)).masked_fill(attn_mask ==0,float(0.0))else:
            attn_mask =None
        self.register_buffer("attn_mask", attn_mask)defforward(self, x):
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W,"input feature has wrong size"
        shortcut = x
        x = self.norm1(x)
        x = x.view(B, H, W, C)'''
          在 H 所在维度上滚动-self.shift_size个像素，
        在 W 所在维度滚动-self.shift_size个像素
        注：torch.roll正方向是从上往下和从左往右
        '''if self.shift_size >0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size,-self.shift_size), dims=(1,2))else:
            shifted_x = x
        '''
        划分窗口
        '''
        x_windows = window_partition(shifted_x, self.window_size)# nW*B, window_size, window_size, C
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)# nW*B, window_size*window_size, C'''
        每个window进行MSA计算
        '''
        attn_windows = self.attn(x_windows, mask=self.attn_mask)# nW*B, window_size*window_size, C'''
        合并窗口
        '''
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, H, W)# B H' W' C'''
        如果是SW-MSA还需要将窗口滑动回来
        '''if self.shift_size >0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1,2))else:
            x = shifted_x
        x = x.view(B, H * W, C)'''
        残差结构
        '''
        x = shortcut + self.drop_path(x)
        x = x + self.drop_path(self.mlp(self.norm2(x)))return x

PatchMerging

论文阅读笔记：Swin Transformer已对PatchMerging进行了介绍，其流程图如下（最终通道翻倍，大小变为原来的一半达到降采样的效果）：
在这里插入图片描述

PatchMerging

的代码和注释如下：

classPatchMerging(nn.Module):""" Patch Merging Layer.
    Args:
        input_resolution (tuple[int]): 输入特征图的大小
        dim (int): 输入特征图的通道
        norm_layer (nn.Module, optional): 归一化层.  Default: nn.LayerNorm
    """def__init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        '''
        将输出通道在最初的特征图通道上翻倍(即流程图中最4个步骤和第1个步骤的倍数关系)
        '''
        self.reduction = nn.Linear(4* dim,2* dim, bias=False)
        self.norm = norm_layer(4* dim)defforward(self, x):"""
        x: B, H*W, C
        """
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W,"input feature has wrong size"assert H %2==0and W %2==0,f"x size ({H}*{W}) are not even."
        x = x.view(B, H, W, C)'''
        将输入按流程图中的方式拆分成4份
        在concat以后进行flatten
        '''
        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
        x = x.view(B,-1,4* C)# B H/2*W/2 4*C'''
        进行LayerNorm并再用一个全连接层输出
        '''
        x = self.norm(x)
        x = self.reduction(x)return x

标签： transformer 深度学习人工智能

本文转载自: https://blog.csdn.net/Z960515/article/details/122784141
版权归原作者 HollowKnightZ 所有，如有侵权，请联系我们删除。

Swin Transformer代码阅读注释

Swin Transformer代码阅读注释

前言

Swin Transformer

1 PatchEmbed

BasicLayer

SwinTransformerBlock

Mlp

window_partition

window_reverse

WindowAttention

Block

PatchMerging

发表评论

“Swin Transformer代码阅读注释”的评论:

关于作者

overfit同步小助手

相关阅读

文章导航