🚩🚩🚩Transformer实战-系列教程总目录
有任何问题欢迎在下面留言
本篇文章的代码运行界面均在Pycharm中进行
本篇文章配套的代码资源已经上传
点我下载源码
SwinTransformer 算法原理
SwinTransformer 源码解读1(项目配置/SwinTransformer类)
SwinTransformer 源码解读2(PatchEmbed类/BasicLayer类)
SwinTransformer 源码解读3(SwinTransformerBlock类)
SwinTransformer 源码解读4(WindowAttention类)
SwinTransformer 源码解读5(Mlp类/PatchMerging类)
5、SwinTransformerBlock类
class SwinTransformerBlock(nn.Module):def extra_repr(self) -> str:return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
5.1 构造函数
SwinTransformerBlock
是 Swin Transformer 模型中的一个基本构建块。它结合了自注意力机制和多层感知机(MLP),并通过窗口划分和可选的窗口位移来实现局部注意力
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_headsself.window_size = window_sizeself.shift_size = shift_sizeself.mlp_ratio = mlp_ratioif min(self.input_resolution) <= self.window_size:self.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)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)if self.shift_size > 0: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)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 = Noneself.register_buffer("attn_mask", attn_mask)
dim
:输入特征的通道数。input_resolution
:输入特征的分辨率(高度和宽度)num_heads
:自注意力头的数量window_size
:窗口大小,决定了注意力机制的局部范围shift_size
:窗口位移的大小,用于实现错位窗口多头自注意力(SW-MSA)mlp_ratio
:MLP隐层大小与输入通道数的比率qkv_bias
:QKV的偏置qk_scale
:QKV的缩放因子drop
:丢弃率drop_path
:分别控制QKV的偏差、缩放因子、丢弃率、注意力丢弃率和随机深度率norm_layer
:激活层和标准化层,默认分别为 GELU 和 LayerNormWindowAttention
:窗口注意力模块Mlp
:一个包含全连接层、激活函数、Dropout的模块img_mask
:图像掩码,用于生成错位窗口自注意力h_slices
和w_slices
:水平和垂直方向上的切片,用于划分图像掩码cnt
:计数器,标记不同的窗口mask_windows
:图像掩码划分为窗口,并将每个窗口的掩码重塑为一维向量window_partition
attn_mask
:注意力掩码,用于在自注意力计算中排除窗口外的位置register_buffer
:注意力掩码注册为一个模型的缓冲区
5.2 前向传播
def forward(self, x):H, W = self.input_resolutionB, L, C = x.shapeassert L == H * W, "input feature has wrong size"shortcut = xx = self.norm1(x)x = x.view(B, H, W, C)if self.shift_size > 0:shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))else:shifted_x = xx_windows = window_partition(shifted_x, self.window_size)x_windows = x_windows.view(-1, self.window_size * self.window_size, C)attn_windows = self.attn(x_windows, mask=self.attn_mask)attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)shifted_x = window_reverse(attn_windows, self.window_size, H, W)if 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)x = shortcut + self.drop_path(x)x = x + self.drop_path(self.mlp(self.norm2(x)))return x
- 原始输入: torch.Size([4, 3136, 96]),输入的是一个长度为3136的序列,每个向量的维度为96,在
被多次调用的时候,维度也发生了变化原始输入: torch.Size([4, 784, 192])、torch.Size([4, 196, 384])、torch.Size([4, 49, 768]) - H,W=[ 56,56],输入分辨率中的高度和宽度
- B, L, C=[ 4,3136,96],当前输入的维度,批次大小、序列长度和向量的维度
- shortcut是用于残差
- norm1(x): torch.Size([4, 3136, 96]),经过一个层归一化,维度不变
- x.view(B, H, W, C): torch.Size([4, 56, 56, 96]),将序列重塑为(Batch_size,Height,Width,Channel)的形状
- shift_size用来判断是否做偏移,最开始的时候shift_size为0,window_partition用来做windows的划分,特征图为[56,56,96],默认窗口为77的大小,因为可以划分出88个窗口
- shifted_x: torch.Size([4, 56, 56, 96]),位移操作后的x,此处的偏移只用了torch的内置函数就可以完成,torch.roll()
- x_windows: torch.Size([256, 7, 7, 96]),使用窗口划分函数,划分256个窗口,每个窗口7*7个特征,每个特征96维向量
- x_windows: torch.Size([256, 49, 96]),将窗口重塑为一维向量,以便进行自注意力计算
- attn_windows: torch.Size([256, 7, 7, 96]),WindowAttention处理x,对每个窗口应用窗口注意力机制,考虑到可能的注意力掩码
- shifted_x: torch.Size([4, 56, 56, 96]),注意力操作后的窗口重塑回原始形状,并将它们合并回完整的特征图
- torch.Size([4, 56, 56, 96]),如果进行了循环位移,则执行逆向循环位移操作,以恢复原始特征图的位置
- torch.Size([4, 3136, 96]),特征图重塑回原始的[B, L, C]形状
- torch.Size([4, 3136, 96]),应用残差连接,并通过随机深度(如果设置了的话)
- torch.Size([4, 3136, 96]),应用第二个标准化层,然后是MLP,并再次应用随机深度,完成残差连接的最后一步。
5.3 窗口划分函数window_partition()
def window_partition(x, window_size):B, H, W, C = x.shapex = 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
在第1个stage中,不同的stage中会做下采样操作
- B, H, W, C =torch.Size([4, 56, 56, 96]),这是函数的原始输入,56*56等于原来的3136
- x = torch.Size([4, 8, 7, 8, 7, 96])
- windows = torch.Size([256, 7, 7, 96]),256的计算方法为原始数据为5656,要分出77大小为一个窗口,所以窗口数量为(56/7)(56/7)=88=64,而64需要乘上batch=4,所以就是256,向量维度96不变
而Windows Attention就是在这7*7里面的窗口进行计算
SwinTransformer 算法原理
SwinTransformer 源码解读1(项目配置/SwinTransformer类)
SwinTransformer 源码解读2(PatchEmbed类/BasicLayer类)
SwinTransformer 源码解读3(SwinTransformerBlock类)
SwinTransformer 源码解读4(WindowAttention类)
SwinTransformer 源码解读5(Mlp类/PatchMerging类)