Swin Transformer源码分析

swin transformer是什么这里就不在说明了，会点进来肯定是知道这个模型是做什么的。

直接看论文有些地方看的一知半解。这里直接从源码分析看下模型的具体实现

论文地址：https://arxiv.org/pdf/2103.14030v1.pdfhttps://arxiv.org/pdf/2103.14030v1.pdf

代码地址：https://github.com/microsoft/Swin-Transformerhttps://github.com/microsoft/Swin-Transformer

首先我们先看下模型结构对下面分析源码很有用处

首先图片 H * W * 3 经过一个 patch partition 缩小四倍同道从 C 变成 48

后续在经过4个 Swin Transformer Block 但是这四个block由

SW-MSA -- W-MSA 循环构成从b图就可以看出来

W-MSA 全称 window multi head self attention

SW-MSA 全称 shift window multi head self attention

下面进入源码分析

class SwinTransformer(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):super().__init__()self.num_classes = num_classesself.num_layers = len(depths)self.embed_dim = embed_dimself.ape = apeself.patch_norm = patch_normself.num_features = int(embed_dim * 2 ** (self.num_layers - 1))self.mlp_ratio = mlp_ratio# split image into non-overlapping patches# 就是模型结构的 Patch Partition 图片变成 (B, H//4 * W//4, embed_dim)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 else None)# 图像缩小4倍后的 patchers numnum_patches = self.patch_embed.num_patches# 图像缩小4倍后的尺寸patches_resolution = self.patch_embed.patches_resolutionself.patches_resolution = patches_resolution# absolute position embeddingif self.ape:# 生成绝对位置编码 num_patches = H//4 * W//4 和经过PatchEmbed后的图片尺寸一致# 对应网络结构中的 linear embedding 网络结构self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))# 绝对位置编码参数初始化trunc_normal_(self.absolute_pos_embed, std=.02)# 添加dropoutself.pos_drop = nn.Dropout(p=drop_rate)# stochastic depth# 给网络层数每层设置随机dropout ratedpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule# build layersself.layers = nn.ModuleList()# 构建四层 w-msa 网络结构# input_resolution 表示每层会缩小 2**i_layer 倍 与给出的模型结构图展示的图像大小缩小倍数对应# depth block 深度# num_heads 多头数量# window_size 窗口大小# mlp_ratio Ratio of mlp hidden dim to embedding dim.# drop_path  dropout rate# downsample 下采样 前三个block 会进行下采样 第四个block 不会在进行下采样for i_layer in range(self.num_layers):layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),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,downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,use_checkpoint=use_checkpoint)self.layers.append(layer)# 层归一化self.norm = norm_layer(self.num_features)# 平均池化self.avgpool = nn.AdaptiveAvgPool1d(1)# 网络输出self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()# 模型所以参数进行初始化self.apply(self._init_weights)def _init_weights(self, m):if isinstance(m, nn.Linear):trunc_normal_(m.weight, std=.02)if isinstance(m, nn.Linear) and m.bias is not None:nn.init.constant_(m.bias, 0)elif isinstance(m, nn.LayerNorm):nn.init.constant_(m.bias, 0)nn.init.constant_(m.weight, 1.0)@torch.jit.ignoredef no_weight_decay(self):return {'absolute_pos_embed'}@torch.jit.ignoredef no_weight_decay_keywords(self):return {'relative_position_bias_table'}def forward_features(self, x):# path_embed 就是模型结构的 Patch Partition 图片变成 (B, H//4 * W//4, embed_dim)x = self.patch_embed(x)# 是否使用绝对位置编码if self.ape:x = x + self.absolute_pos_embedx = self.pos_drop(x)# 经过4个 swin transformer blockfor 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 xdef forward(self, x):# 进行前向计算x = self.forward_features(x)# 模型最后输出用来进行分类 (b, c)x = self.head(x)return xdef flops(self):"""这个方法是用来计算模型性能的floating point operations per second"""flops = 0flops += self.patch_embed.flops()for i, layer in enumerate(self.layers):flops += layer.flops()flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)flops += self.num_features * self.num_classesreturn flops

从上面我们可以看出网络结构和论文中给出的图基本差不多但是从代码看图网络结构更像是下面红色划分

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

前三个block 会进行下采样 第四个block 不会在进行下采样  看上去更像是

4个（block + patch merging）

然后分析下

PatchEmbed 和 PatchMerging

class PatchEmbed(nn.Module):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_sizeself.patch_size = patch_sizeself.patches_resolution = patches_resolutionself.num_patches = patches_resolution[0] * patches_resolution[1]self.in_chans = in_chansself.embed_dim = embed_dim# 用一个卷积操作来实现图像缩小四倍 kernel_size = 4  stride = patch_sizeself.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 = Nonedef forward(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]})."# shape 变化 B, C, H, W --> B, C, h, w --> B, C, h * w --> B, h * w, cx = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C# 进行层归一化if self.norm is not None:x = self.norm(x)return x

实现就是用一个 kernel_size = 4 stride = patch_size 的卷积操作来实现

nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

class PatchMerging(nn.Module):def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):super().__init__()self.input_resolution = input_resolutionself.dim = dimself.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_resolutionB, L, C = x.shapeassert 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)# 这里实现path merging  图片缩小一半# 这里解释下 # 0::2 从 0 开始 隔一个点取一个值 # 1::2 从 1 开始 隔一个点取一个值x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 Cx1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 Cx2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 Cx3 = 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*Cx = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C# 层归一化x = self.norm(x)# 降维到2 * dim 图片缩小一倍 通道维度增加一倍x = self.reduction(x)return x

可以从下面的表格看出 x0 x1 x2 x3 分别对应表格中 0 1 2 3 对应位置的点最后在通道上合并图片缩小一倍。

0	2	0	2
1	3	1	3
0	2	0	2
1	3	1	3

这里我们可以想一下实现这种patch merge 方法有很多比如用2*2卷积来实现、2*2平均池化等。模型性能会提升还是降低？

接下来分析basic layer 层

class BasicLayer(nn.Module):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# 当前层有多少个SwinTransformerBlockself.depth = depthself.use_checkpoint = use_checkpoint# build blocks# 构建深度为depth的block堆叠self.blocks = nn.ModuleList([# shift_size 需要注意下# 偶数个block是进行W-MSA 奇数进行SW-MSA   # SW-MSA - W-MSA - SW-MSA - W-MSA 的循环# 这样让输出特征包含 local window attention 和 跨窗口的 window attentionSwinTransformerBlock(dim=dim, input_resolution=input_resolution,num_heads=num_heads, window_size=window_size,shift_size=0 if (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] if isinstance(drop_path, list) else drop_path,norm_layer=norm_layer)for i in range(depth)])# patch merging layer# 是否进行 patch merging  前三个block 执行patch merging # 最后一个block 不会执行 patch merging if downsample is not None:self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)else:self.downsample = Nonedef forward(self, x):# 进行前向传播for blk in self.blocks:if self.use_checkpoint:x = checkpoint.checkpoint(blk, x)else:x = blk(x)# 是否进行 patch merging if self.downsample is not None:x = self.downsample(x)return x

进入 SwinTransformerBlock 我们主要分析 SW-MSA 的实现 W-MSA的实现很简单就是简单的局部 window multi head self attention 熟悉 Bert 和 transformer模型的人肯定很清楚 qkv三个矩阵的计算公式。

class SwinTransformerBlock(nn.Module):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 = dimself.input_resolution = input_resolutionself.num_heads = num_heads# 默认大小 7self.window_size = window_size# 进行 SW-MSA shift-size 7//2=3# 进行 W-MSA shift-size 0self.shift_size = shift_size# multi self attention 最后神经网络的隐藏层的维度self.mlp_ratio = mlp_ratioif min(self.input_resolution) <= self.window_size:# if window size is larger than input resolution, we don't partition windows# 简单的判定 如果 最后图像缩小到比window size 还小 调整 window size 大小# 将 shift_size 赋值为 0  也就是说直接进行 W-MSAself.shift_size = 0self.window_size = min(self.input_resolution)assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"# 层归一化self.norm1 = norm_layer(dim)# local window multi head self attentionself.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)# dropout rateself.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)# 可以看出 上面 结构是  layer normal +  W-MSA/SW-MSA + layer normal + mlpif self.shift_size > 0:# nW * B, window_size * window_size, C# calculate attention mask for SW-MSA# attention mask 的构成H, W = self.input_resolutionimg_mask = torch.zeros((1, H, W, 1))  # 1 H W 1h_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, window_size, window_size, 1mask_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 = Noneself.register_buffer("attn_mask", attn_mask)def forward(self, x):H, W = self.input_resolutionB, L, C = x.shapeassert L == H * W, "input feature has wrong size"shortcut = x# 层归一化 如图所示  layer normal + W-MSA/SW-MSAx = self.norm1(x)x = x.view(B, H, W, C)# cyclic shiftif self.shift_size > 0:# 进行 sw-msa 将数据进行变换shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))else:shifted_x = x# partition windows# 将数据拆分成  n * window_size * window_size * c 的维度 方便进行self attentionx_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C# 数据reshape成 (n, window_size * window_size, c) 送入 attention 层x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C# W-MSA/SW-MSA# 进行 multi head attentionattn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C# merge windowsattn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)# 将数据维度退回到window_partition之前的维度shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C# reverse cyclic shiftif self.shift_size > 0:x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))else:x = shifted_xx = x.view(B, H * W, C)# FFNx = shortcut + self.drop_path(x)x = x + self.drop_path(self.mlp(self.norm2(x)))return x

在分析代码前，我们先看下作者为什么会设计SW-MSA？

论文所述：

The shifted windows bridge the windows of the preceding layer, providing connections among them that signifificantly enhance modeling power

简单来说，就是SW-MSA 是为了建立两个widow之间的桥梁设计的。也就是融合两个window之间的特征。极大的加强了模型的能力。

Wait. 这不就是类似滑动卷积核做卷积？真有你的哦！

接下来就是正题，源码中最难理解的地方来了。这里我先提前说下，作者是通过设计一个MASK来实现SW-MSA。

下面我们用一段代码模拟下：

import torchdef window_partition(x, window_size):H, W = x.shapex = x.view(H // window_size, window_size, W // window_size, window_size)windows = x.permute(0, 2, 1, 3).contiguous().view(-1, window_size, window_size)return windowswindow_size = 3
shift_size = 3 // 2data = torch.arange(81).view(9, 9)shift_data = torch.roll(data, shifts=(-shift_size, -shift_size), dims=(0, 1))mask = torch.zeros(9, 9)h_slices = (slice(0, -window_size),slice(-window_size, -shift_size),slice(-shift_size, None))w_slices = (slice(0, -window_size),slice(-window_size, -shift_size),slice(-shift_size, None))cnt = 0
for h in h_slices:for w in w_slices:mask[h, w] = cntcnt += 1print('data', data)
print('shift_data', shift_data)
print('mask', mask)mask_windows = window_partition(mask, window_size)  # nW, window_size, window_size
print('mask_windows', mask_windows)mask_windows = mask_windows.view(-1, window_size * window_size)print('reshape_mask_windows', mask_windows)attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)print('attn_mask', attn_mask)attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))print('fill_attn_mask', attn_mask)

data 输出

tensor([[ 0, 1, 2, 3, 4, 5, 6, 7, 8],
[ 9, 10, 11, 12, 13, 14, 15, 16, 17],
[18, 19, 20, 21, 22, 23, 24, 25, 26],
[27, 28, 29, 30, 31, 32, 33, 34, 35],
[36, 37, 38, 39, 40, 41, 42, 43, 44],
[45, 46, 47, 48, 49, 50, 51, 52, 53],
[54, 55, 56, 57, 58, 59, 60, 61, 62],
[63, 64, 65, 66, 67, 68, 69, 70, 71],
[72, 73, 74, 75, 76, 77, 78, 79, 80]])

shift_data 输出

tensor([[10, 11, 12, 13, 14, 15, 16, 17, 9],
[19, 20, 21, 22, 23, 24, 25, 26, 18],
[28, 29, 30, 31, 32, 33, 34, 35, 27],
[37, 38, 39, 40, 41, 42, 43, 44, 36],
[46, 47, 48, 49, 50, 51, 52, 53, 45],
[55, 56, 57, 58, 59, 60, 61, 62, 54],
[64, 65, 66, 67, 68, 69, 70, 71, 63],
[73, 74, 75, 76, 77, 78, 79, 80, 72],
[ 1, 2, 3, 4, 5, 6, 7, 8, 0]])

上面两个输出刚好对应下图的 cyclic shift

经过cyclic shift 后的数据被分成了9份如下图所示：每份之间的数据是互相可见的，其中1单独组成个window [2，3] 组成一个window 且 [2，3] 之间的数据互相不可见，但是 [2，3] 内的数据互相可见。同理 [4，7] 组成一个widnow。 [5, 6, 8, 9] 组成一个window 这是最特殊的一个window 由三部分shift 出去的数据和原先最后一个widnow剩下的数据组成。它们之间数据的可见性同上。

现在再看下面代码是不是瞬间理解了。

将mask分成上面的9份分别用0~8设置

mask = torch.zeros(9, 9)
h_slices = (slice(0, -window_size),slice(-window_size, -shift_size),slice(-shift_size, None))
w_slices = (slice(0, -window_size),slice(-window_size, -shift_size),slice(-shift_size, None))
cnt = 0
for h in h_slices:for w in w_slices:mask[h, w] = cntcnt += 1

mask 输出就和上面分析的一摸一样

tensor([[0., 0., 0., 0., 0., 0., 1., 1., 2.],
[0., 0., 0., 0., 0., 0., 1., 1., 2.],
[0., 0., 0., 0., 0., 0., 1., 1., 2.],
[0., 0., 0., 0., 0., 0., 1., 1., 2.],
[0., 0., 0., 0., 0., 0., 1., 1., 2.],
[0., 0., 0., 0., 0., 0., 1., 1., 2.],
[3., 3., 3., 3., 3., 3., 4., 4., 5.],
[3., 3., 3., 3., 3., 3., 4., 4., 5.],
[6., 6., 6., 6., 6., 6., 7., 7., 8.]])

window_partition 就是很简单的 reshape 操作将器拆分成一个个window_size * window_size大小的window。从 (9, 9) - > (9, 3, 3)

def window_partition(x, window_size):H, W = x.shapex = x.view(H // window_size, window_size, W // window_size, window_size)windows = x.permute(0, 2, 1, 3).contiguous().view(-1, window_size, window_size)return windowsmask_windows = window_partition(mask, window_size)
mask_windows = mask_windows.view(-1, window_size * window_size)

mask_window 输出 9 * 9 每一行代表一个window共9个window 每一列代表一个window内mask的值 window 大小 window_size * window_size （window_size = 3）

tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[1., 1., 2., 1., 1., 2., 1., 1., 2.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[1., 1., 2., 1., 1., 2., 1., 1., 2.],
[3., 3., 3., 3., 3., 3., 6., 6., 6.],
[3., 3., 3., 3., 3., 3., 6., 6., 6.],
[4., 4., 5., 4., 4., 5., 7., 7., 8.]])

attn_mask 很多人想不明白这里是在干什么？没关系往后看，我们一步一步分析

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

mask_windows shape: (nw, 3 * 3) nw: widown数量 3: 窗口大小

a = mask_windows.unsqueeze(1) shape: (nw, 1, 3* 3)

b = mask_windows.unsqueeze(2) shape: (nw, 3* 3 , 1)

由于广播机制 a, b 向对方的维度扩展变成 (nw, 3* 3, 3* 3)

其中 a b 在 dim = [1, 2] 维度上互为转置矩阵 a - b 类似于行减去列的值

此时每行的值是同一个window的的所有mask值，每一列的值代表当前window里第i个mask值

而mask值由上面分析是由 0~8之间的数字组成，相同的表示互相可见。

如果两者相互可见，表示mask值一样，相减等于0。

此时 c = a - b shape (nw, 3 * 3, 3 * 3)

c[i, j, k] 表示第 i 个window 内第 j 个值与 window 内第 k 个值是否互相可见 0 < j, k < 3 * 3

是不是感觉到熟悉了，这不就和self attention里面的 q * k.T后表示的意思一样嘛。

下面就是将非0值用一个极大的负值替换。用来当 SW-MSA 的 mask 值。因为 self attention中是用softmax 来得到加权值所以用一个大的负数来填充，得到的softmax值越接近0

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

最后我们来看下 WindowAttention的实现。这里的实现只需要注意两点，就是pos编码设计和transform里面不同。

1. 设计了个相对位置编码bias表格

class WindowAttention(nn.Module):def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):super().__init__()self.dim = dimself.window_size = window_size  # Wh, Wwself.num_heads = num_headshead_dim = dim // num_headsself.scale = qk_scale or head_dim ** -0.5# define a parameter table of relative position bias# 设计了一个相对位置编码  相对位置编码个数 (2 * window_size - 1) * (2 * window_size - 1)# 因为在一个 7 * 7 的 window内  相对位置范围 (-6, 6) 有 2 * 7 - 1 = 13 个# 所以对于二维数据 应该有 13 * 13 个 即 (2w - 1) * (2w - 1) 也就是相对坐标范围应该为 0 ~ (2w - 1) * (2w - 1) - 1self.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# get pair-wise relative position index for each token inside the window# 得到 window 内的表格坐标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# 下面两行其实和 mask生成的是做的操作类似  也就是将数据展平后 广播后相减 得到相对坐标coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Wwrelative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Wwrelative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2# 但是上面的相对坐标和定义的 relative_position_bias_table还对应不上 relative_coords 取值范围 (-w + 1) ~ (w - 1)# 所以在dim=[1, 2] 维度 才加上 self.window_size[0] - 1 取值范围 变成 0 ~ (2w - 2)relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0relative_coords[:, :, 1] += self.window_size[1] - 1# 最后在 dim=2 的维度上 乘以 2w - 1 所以在这个维度上取值范围为 0 ~  (2w - 2) * (2w - 1) = (2w - 1)**2 - (2w - 1)relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1# 最后求和后得到的相对坐标范围 0 ~ (2w - 1)**2 - (2w - 1) + (2w - 2) = (2w - 1)**2 - 1# OKay 到此为止终于得到范围为 0 ~ (2w - 1)**2 - 1 和 上面的 relative_position_bias_table对应上了# 所以每次只需要用相对索引去 relative_position_bias_table 表格中取值就行了relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Wwself.register_buffer("relative_position_index", relative_position_index)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)

2. 添加位置编码信息的位置不同，这里是在计算 qk.T后在加上位置编码 bias

attn = (q @ k.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)

其他的就和普通的 self attention 差不多这里就不进行分析了。

到此为止源码基本分析完！

从代码上我们可以看到，Swin Transformer 基本上是抛弃了卷积操作。但一个SW-MSA确又看到了卷积的影子。

SwinTransformV2 已经推出，用于大模型。对网络结构也做了一些调整。

后续我也会对SwinTransformV2源码进行分析。

Swin Transformer源码分析相关推荐

DataX Transformer 源码分析及 UDF 扩展与使用
DataX GitHub DataX Transformer 目录 1 前言 2 需求说明 3 解决方案分析 4 解密算法 5 Hive UDF 5.1 测试数据 5.2 新建 Maven 项目 5. ...
DataX Transformer从入口到加载的源码分析及UDF扩展与使用
DataX GitHub DataX Transformer 目录 1 前言 2 需求说明 3 解决方案分析 4 解密算法 5 Hive UDF 5.1 测试数据 5.2 新建 Maven 项目 5. ...
Vuex 2.0 源码分析
作者:滴滴公共前端团队 - 黄轶大家好,我叫黄轶,来自滴滴公共前端团队,我们团队最近写了一本书 --<Vue.js 权威指南>,内容丰富,由浅入深.不过有一些同学反馈说缺少 Vuex 的 ...
BERT源码分析（PART III）
写在前面继续之前没有介绍完的 Pre-training 部分,在上一篇中(BERT源码分析(PART II))我们已经完成了对输入数据的处理,接下来看看 BERT 是怎么完成「Masked LM」和 ...
vuex 源码分析_Vuex 2.0 源码分析（下）
大家好,我叫黄轶,来自滴滴公共前端团队,最近在幕课网上线了一门 Vue.js 的实战课程--<Vue.js高仿饿了么外卖App 2016最火前端框架>,同时,我们团队最近写了一本书 --& ...
图片加载框架Picasso －源码分析
简书:图片加载框架Picasso - 源码分析前一篇文章讲了Picasso的详细用法,Picasso 是一个强大的图片加载缓存框架,一个非常优秀的开源库,学习一个优秀的开源库,,我们不仅仅是学习它的 ...
bytebuddy实现原理分析源码分析（三）- advice 详解
advice详解八.advice 8.1 AsmVisitorWrapper 8.1.1 ForDeclareFields 8.1.1.1 Entry 8.1.1.2 DispatchingVisi ...
10章 RxJava源码分析
本篇文章已授权微信公众号 YYGeeker 独家发布转载请标明出处 CSDN学院课程地址 RxJava2从入门到精通-初级篇:https://edu.csdn.net/course/detail/10 ...
bytebuddy实现原理分析源码分析（二）
bytebuddy实现原理分析 &源码分析四.字节码的操作(implementation ) pkg 4.1 bytecode :pkg 4.1.1 StackManipulation :c ...

Swin Transformer源码分析

Swin Transformer源码分析相关推荐

最新文章

热门文章