【神经网络架构】Swin Transformer细节详解-1

目录

源码：

1. Patch Partition + Liner Embedding 模块

2. Swin Transformer block(一个完整的W-MSA)

partition windows

 W-MSA

相对位置偏差

 merge windows

MLP

补充(熟悉的话直接看 3. Patch Merging)

3. Patch Merging(downsample)

4. 图像分类任务

5. 待续

6. 参考

---

图3。(a) Swin Transformer (Swin-T)的架构;
(b)两个连续的Swin Transformer Blocks(符号见公式(3))。W-MSA和SW-MSA分别是具有规则（regular）窗口配置和移动（shifted）窗口配置的multi-head self attention modules。

源码：

GitHub - microsoft/Swin-TransformerThis is an official implementation for "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows". - GitHub - microsoft/Swin-Transformer: This is an official implementation for "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows".https://github.com/microsoft/Swin-Transformer

1. Patch Partition + Liner Embedding 模块

代码：

class PatchEmbed(nn.Module):
    r""" Image to Patch Embedding

    Args:
        img_size (int): Image size.  Default: 224.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
        super().__init__()
        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 = in_chans
        self.embed_dim = embed_dim

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert 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]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
        if self.norm is not None:
            x = self.norm(x)
        return x

    def flops(self):
        Ho, Wo = self.patches_resolution
        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
        if self.norm is not None:
            flops += Ho * Wo * self.embed_dim
        return flops

这里，Patch Partition + Liner Embedding的实现是通过nn.Conv2d，将kernel_size和stride设置为 patch_size（将图像分为几块，这里为4）大小。其实就是一个2d卷积实现的滑动窗口。然后，为了在通道维度C（代码中为96）进行LayerNorm，进行了flatten(2).transpose(1, 2)  # （B h/4 * w/4 C）。 上图中的N即批量大小B。 

 实现：

self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) # 实现4块划分，并且将dim 转为embed_dim
if norm_layer is not None:
    self.norm = norm_layer(embed_dim)
else:
    self.norm = None


B, C, H, W = x.shape # 输入image的size
x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C 因为需要对C做LN
if self.norm is not None:
    x = self.norm(x)

2. Swin Transformer block(一个完整的W-MSA)

        先不考虑cyclic shift。

partition windows

window partition函数是用于按照指定窗口大小对张量划分窗口。将原本的张量从 B H W C, 划分成 num_windows*B, window_size, window_size, C，其中 num_windows = H*W / (window_size*window_size)，即窗口的个数。

代码：

def window_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
    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

x_windows: (nW*B, window_size, window_size, C )    (nW*B, window_size*window_size, C)

 W-MSA

# W-MSA/SW-MSA
attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C

这里的mask表示是否需要在计算self-attention时掩码。下面讨论不带mask的MSA。

self.attn = WindowAttention(
            dim,                                      # 输入通道的个数
            window_size=to_2tuple(self.window_size),  # 变成两个元组，比如7变成（7，7）
            num_heads=num_heads,                      # 注意力头个数
            qkv_bias=qkv_bias,                        # (bool, 可选): If True, add a learnable bias to query, key, value. Default: True
            qk_scale=qk_scale,                        # (float | None, 可选): Override default qk scale of head_dim ** -0.5 if set. 
            attn_drop=attn_drop,                      # Attention dropout rate. Default: 0.0
            proj_drop=drop)                           # Stochastic depth rate. Default: 0.0

class WindowAttention(nn.Module)中

def forward(self, x, mask=None):
    """
    Args:
        x: input features with shape of (num_windows*B, N, C) （nW*B, window_size*window_size, C）
        mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
    """
    B_, N, C = x.shape # （nW*B, window_size*window_size, C）
    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))

self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # qkv_bias=True

 (nW*B, window_size*window_size, C)  (nW*B, window_size*window_size, 3C)

(nW*B, window_size*window_size, 3, num_heads, C // num_heads)qkv: (3, nW*B, num_heads, window_size*window_size, C // num_heads)

这里阶段1，2，3，4的Swin Transformer block的 num_heads分别为[3, 6, 12, 24]。这里C在每个Swin Transformer block中都会加倍，而num_heads也加倍。故q, k, v 的 C // num_heads为固定值。假设Patch Partition + Liner Embedding 模块的输出中C 为 96， 则 C // num_heads固定为32。

q,k,v的维度(nW*B, num_heads, window_size*window_size, C // num_heads)。

在每个划出来的窗口中做 self-attention。例如，窗口为7，则一共 7*7 = 49个元素，每个元素会做49次 self-attention ，因此输出为 (49 ，49)大小。

head_dim = dim // num_heads

self.scale = qk_scale or head_dim ** -0.5 # qk_scale = None ,  1/sqrt(d)

q:  (nW*B,  num_heads,  window_size*window_size,  C // num_heads) *  1/sqrt(head_dim)

k.transpose(-2, -1)： (nW*B, num_heads, C // num_heads，window_size*window_size)

 = attn: (nW*B,  num_heads,  window_size*window_size,  window_size*window_size)

相对位置偏差

         上图中，红色为相对位置偏差。论文中，Q,K,V维度：(window_size*window_size,  C // num_heads)。和B的维度为：(window_size*window_size, window_size*window_size)。

这里的例子为M = window_size = 2。相对位置偏差 B 是通过 相对位置索引 查 相对位置偏差表(论文中的 ) 得到的。相对位置索引是根据M大小固定的，相对位置偏差表是训练得到的。 

那么如何得到相对位置索引？

假设M = window_size = 2。

 

 

 由于最终我们希望使用一维的位置坐标 x+y 代替二维的位置坐标(x,y)，为了避免 (0,-1) (-1,0) 两个坐标转为一维时均为-1，我们之后对相对位置索引进行了一些线性变换，使得能通过一维的位置坐标唯一映射到一个二维的位置坐标。

 为什么相对位置索引有(2M-1) * (2M-1)种状态，以及为什么相对位置偏差表的长度为 (2M-1) * (2M-1)?

        相对位置索引的取值范围为 [-M + 1, M-1]， 那么坐标(行, 列)取值有(2M-1)种状态。如M = 2，则坐标(行, 列)取值范围都为 [-1, 1]，有3种状态，所以一共为9种状态。

        所以相对位置偏差表的长度为 (2M-1) * (2M-1) 。

实现：

window_size = (2, 2)
# get pair-wise relative position index for each token inside the window
coords_h = torch.arange(window_size[0])
coords_w = torch.arange(window_size[1])
coords = torch.meshgrid([coords_h, coords_w])
# tensor([[0, 0],
#         [1, 1]]), 
# tensor([[0, 1],
#         [0, 1]])
coords = torch.stack(coords)  # 2, Wh, Ww
# tensor([[[0, 0],
#          [1, 1]],

#         [[0, 1],
#          [0, 1]]])
coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
# tensor([[0, 0, 1, 1],
#         [0, 1, 0, 1]])
relative_coords_first = coords_flatten[:, :, None]  # 2, wh*ww, 1
# tensor([[[0],
#          [0],
#          [1],
#          [1]],

#         [[0],
#          [1],
#          [0],
#          [1]]])
relative_coords_second = coords_flatten[:, None, :] # 2, 1, wh*ww
# tensor([[[0, 0, 1, 1]],

#         [[0, 1, 0, 1]]])
relative_coords = relative_coords_first - relative_coords_second  # 2, Wh*Ww, Wh*Ww # 对前者的元素进行广播来相减
# tensor([[[ 0,  0, -1, -1],
#          [ 0,  0, -1, -1],
#          [ 1,  1,  0,  0],
#          [ 1,  1,  0,  0]],  # x (行) 坐标

#         [[ 0, -1,  0, -1],
#          [ 1,  0,  1,  0],
#          [ 0, -1,  0, -1],
#          [ 1,  0,  1,  0]]]) # y(列) 坐标
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2   # 2表示(x, y)坐标 即 (行, 列）
# tensor([[[ 0,  0],   # 每对是一个 (x, y)坐标
#          [ 0, -1],
#          [-1,  0],
#          [-1, -1]],

#         [[ 0,  1],
#          [ 0,  0],
#          [-1,  1],
#          [-1,  0]],

#         [[ 1,  0],
#          [ 1, -1],
#          [ 0,  0],
#          [ 0, -1]],

#         [[ 1,  1],
#          [ 1,  0],
#          [ 0,  1],
#          [ 0,  0]]])
relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0  # * (M-1)
# tensor([[[ 1,  0],      # 所有 x (行) 坐标从0 开始
#          [ 1, -1],
#          [ 0,  0],
#          [ 0, -1]],

#         [[ 1,  1],
#          [ 1,  0],
#          [ 0,  1],
#          [ 0,  0]],

#         [[ 2,  0],
#          [ 2, -1],
#          [ 1,  0],
#          [ 1, -1]],

#         [[ 2,  1],
#          [ 2,  0],
#          [ 1,  1],
#          [ 1,  0]]])
relative_coords[:, :, 1] += window_size[1] - 1  # 所有 y(列) 坐标从0 开始 # * (M-1)
# tensor([[[1, 1],
#          [1, 0],
#          [0, 1],
#          [0, 0]],

#         [[1, 2],
#          [1, 1],
#          [0, 2],
#          [0, 1]],

#         [[2, 1],
#          [2, 0],
#          [1, 1],
#          [1, 0]],

#         [[2, 2],
#          [2, 1],
#          [1, 2],
#          [1, 1]]])
relative_coords[:, :, 0] *= 2 * window_size[1] - 1 # x (行) 坐标 * (2M -1) # Wh*Ww, Wh*Ww, 2 
# tensor([[[3, 1],
#          [3, 0],
#          [0, 1],
#          [0, 0]],

#         [[3, 2],
#          [3, 1],
#          [0, 2],
#          [0, 1]],

#         [[6, 1],
#          [6, 0],
#          [3, 1],
#          [3, 0]],

#         [[6, 2],
#          [6, 1],
#          [3, 2],
#          [3, 1]]])
relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
# tensor([[4, 3, 1, 0],   # (x, y)坐标 -> 求和变成 1-d
#         [5, 4, 2, 1],
#         [7, 6, 4, 3],
#         [8, 7, 5, 4]])
# self.register_buffer("relative_position_index", relative_position_index) # 注册为一个不参与网络学习的变量

relative_position_index = relative_position_index.view(-1)
# tensor([4, 3, 1, 0, 5, 4, 2, 1, 7, 6, 4, 3, 8, 7, 5, 4])

num_heads = 1
relative_position_bias_table = torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads) # 2*Wh-1 * 2*Ww-1, nH
trunc_normal_(relative_position_bias_table, std=.02) # 截断 正态分布 relative_position_bias_table为训练的参数 nn.Parameter
# tensor([[ 0.0121],
#         [-0.0030],
#         [ 0.0043],
#         [ 0.0263],
#         [ 0.0264],
#         [ 0.0187],
#         [ 0.0364],
#         [ 0.0182],
#         [-0.0170]])
relative_position_value = relative_position_bias_table[relative_position_index] # Wh*Ww * Wh*Ww, nH  查表 
relative_position_bias = relative_position_value.view(window_size[0] * window_size[1], window_size[0] * 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
relative_position_bias = relative_position_bias.unsqueeze(0) # 1, nH, Wh*Ww, Wh*Ww

attn = attn + relative_position_bias # 对relative_position_bias广播来逐元素相加 # (nW*B,  num_heads,  window_size*window_size,  window_size*window_size)

# 先不考虑 mask 
self.attn_drop = nn.Dropout(attn_drop) #attn_drop: Dropout ratio of attention weight. Default: 0.0

attn = self.softmax(attn)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
x = self.proj(x) # nn.Linear(dim, dim)
x = self.proj_drop(x) # nn.Dropout(proj_drop) Default: 0.0

attn: (nW*B,  num_heads,  window_size*window_size,  window_size*window_size)v: (nW*B, num_heads, window_size*window_size,  C // num_heads) = (nW*B,  num_heads,  window_size*window_size, C // num_heads)  (nW*B,   window_size*window_size, num_heads,  C // num_heads) (nW*B,   window_size*window_size, C) attn_windows: (nW*B,   window_size*window_size, C)

 merge windows

# merge windows
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  这里第一个Transformer:H=W=56

attn_windows: (nW*B,   window_size*window_size, C) (nW*B,   window_size, window_size, C)  (B, H, W, C)

def window_reverse(windows, window_size, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C) # num_windows*B = H*W / (window_size*window_size) *B
        window_size (int): Window size
        H (int): Height of image  # 224/4 = 56/[1,2,2^2, 2^3]
        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) # (B, H // window_size, W // window_size, window_size, window_size, C)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) # (B, H, W, C)
    return x

先不先考虑reverse cyclic shift，见下面。

(B, H, W, C) x: (B, H*W, C)

# FFN
x = shortcut + self.drop_path(x) # (B, H*W, C)
x = x + self.drop_path(self.mlp(self.norm2(x)))

MLP

mlp_ratio=4.
act_layer=nn.GELU
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

class Mlp(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)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

补充(熟悉的话直接看 3. Patch Merging)

        其中，shortcut：(B, H*W, C)。有个self.drop_path(x)操作。

self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

 torch.nn.identity()方法详解_sigmoidAndReLU的博客-CSDN博客_torch.nn.identity()

nn.Identity() ：用于占位，什么也不做。

DropPath：DropPath - 巴蜀秀才 - 博客园

实现：

def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None, scale_by_keep=True):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

测试：

import torch
x = torch.randn(2, 1, 2, 2)
print(x)
keep_prob = 0.5
shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
print(shape)
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
print(random_tensor)
# random_tensor.div_(keep_prob)
print(x * random_tensor)

输出：

tensor([[[[ 0.3913,  0.4729],
          [ 0.2470, -0.7110]]],

        [[[ 0.2733,  0.6184],
          [-0.2881,  0.3545]]]])

(2, 1, 1, 1)

tensor([[[[1.]]],

        [[[0.]]]])

tensor([[[[ 0.3913,  0.4729],
          [ 0.2470, -0.7110]]],

        [[[ 0.0000,  0.0000],
          [-0.0000,  0.0000]]]])

        DropPath是对Batch = 1, 输出为全 0。若x为输入的张量，其通道为[B,H,W, C]，那么drop_path的含义为在一个Batch_size中，随机有drop_prob的某些样本，直接置0。

3. Patch Merging(downsample)

     

        论文架构图中，Stage2，3，4为Patch Merging + Swin Transformer block。而代码实现中是将Swin Transformer block与Patch Merging组合在一起构造了Class BasicLayer。

而最后一个 BasicLayer_3中没有Patch Merging。代码中self.num_layers = 4，如下：

layer = BasicLayer(..., downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, ...)

       Patch Merging是用来降低分块图像的分辨率x0.5，而特征通道数 x2。这类似于CNN中在每个block之前的用stride=2的卷积/池化层来降低分辨率，同时特征通道加倍。

实现：在行和列的方向上，隔一个取元素；在通道维拼接成一个张量；然后在对张量做变换用来在通道维进行LayerNorm操作。最后使用nn.Linear调整通道维数。

代码：

class PatchMerging(nn.Module):
    r""" Patch Merging Layer.

    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)

    def forward(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 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."

        x = x.view(B, H, W, C)

        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

        x = self.norm(x)
        x = self.reduction(x)

        return x

4. 图像分类任务

self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))  # 96 * 2^3 =768
self.norm = nn.LayerNorm(self.num_features) 
self.avgpool = nn.AdaptiveAvgPool1d(1)
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

x = self.norm(x)  # B L C
x = self.avgpool(x.transpose(1, 2))  # B C 1
x = torch.flatten(x, 1) # B C

x = self.head(x) # B num_classes

        例如，输入为(128, 3, 224, 224)，BasicLayer_3 后的输出为 (128, 49, 768)，代码中C为96。

(128, 49, 768) (128, 49, 768) (128, 768, 49) (128, 768, 1) (128, 768) (128, num_classes)

5. 待续

1. cyclic shift + reverse cyclic shift
2. SW-MSA
3. W-MSA和MSA的复杂度对比

6. 参考

图解Swin Transformer - 知乎

论文详解：Swin Transformer - 知乎

https://github.com/microsoft/Swin-Transformer

12.1 Swin-Transformer网络结构详解_哔哩哔哩_bilibili