Swin Transformer模子详解(附pytorch实现)

写在前面

Swin Transformer（Shifted Window Transformer）是一种新颖的视觉Transformer模子，在2021年由微软亚洲研究院提出。这一模子提出了一种基于局部窗口的自注意力机制，明显改善了Vision Transformer（ViT）在处理高分辨率图像时的性能，尤其是在图像分类、物体检测等计算机视觉任务中表现出色。
Swin Transformer的最大创新之一是其引入了“平移窗口”机制，降服了传统自注意力方法在大图像处理时计算资源斲丧过大的问题。这一机制使得模子可以或许在不同层次上以局部的方式计算自注意力，同时保持全局信息的处理能力。
在本文中，我们将通过具体的分析，介绍Swin Transformer的模子布局、核心头脑及实在现，最后提供一个基于PyTorch的简单实现。
论文地址：https://arxiv.org/pdf/2103.14030
官方代码实现：https://github.com/microsoft/Swin-Transformer
Swin网络布局

如下图所示，Swin Transformer的Encoder接纳分层的方式，通过多个阶段（Stage）渐渐镌汰特性图的分辨率，同时增加特性维度。每个Stage包罗多少个Transformer Block。

每个Block通常由以下几个部分组成：

• Window-based Self-Attention：每个Block使用窗口自注意力机制，在每个窗口内计算自注意力。这种方式镌汰了计算量，因为自注意力只在局部窗口内进行计算，而不是整个图像。
  
• Shifted Window：为了增强不同窗口之间的接洽，Swin Transformer在每一层的Block中接纳了“窗口位移”策略。每一层中的窗口会偏移肯定的步长，使得窗口之间的重叠地区增加，从而促进信息交换。
  

Patch Partition

Patch Partition 是将输入图像分割成固定大小的块（patch）并将其映射到高维空间的操作。就相称于是VIT模子当中的 Patch Embedding。

1. from functools import partial
3. import torch
4. import torch.nn as nn
5. import torch.nn.functional as F
6. from pyzjr.utils.FormatConver import to_2tuple
7. from pyzjr.nn.models.bricks.drop import DropPath
9. LayerNorm = partial(nn.LayerNorm, eps=1e-6)
12. class PatchPartition(nn.Module):
13.     def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None):
14.         super().__init__()
15.         self.patch_size = to_2tuple(patch_size)
16.         self.embed_dim = embed_dim
17.         self.proj = nn.Conv2d(in_channels, self.embed_dim,
18.                               kernel_size=self.patch_size, stride=self.patch_size)
19.         self.norm = norm_layer(self.embed_dim) if norm_layer else nn.Identity()
21.     def forward(self, x):
22.         _, _, H, W = x.shape
23.         if H % self.patch_size[0] != 0:
24.             pad_h = self.patch_size[0] - H % self.patch_size[0]
25.             x = F.pad(x, (0, 0, 0, pad_h))
26.         if W % self.patch_size[1] != 0:
27.             pad_w = self.patch_size[1] - W % self.patch_size[1]
28.             x = F.pad(x, (0, pad_w, 0, 0))
29.         x = self.proj(x)     # [B, embed_dim, H/patch_size, W/patch_size]
30.         Wh, Ww = x.shape[2:]
31.         x = x.flatten(2).transpose(1, 2)    # [B, num_patches, embed_dim]
32.         # Linear Embedding
33.         x = self.norm(x)
34.         # x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
35.         return x, Wh, Ww
38. if __name__=="__main__":
39.     batch_size = 1
40.     in_channels = 3
41.     height, width = 30, 32
42.     patch_size = 4
43.     embed_dim = 96
45.     x = torch.randn(batch_size, in_channels, height, width)
46.     patch_partition = PatchPartition(patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim)
47.     output,_ ,_ = patch_partition(x)
49.     print(f"Output shape: {output.shape}")

复制代码

Patch Merging

PatchMerging 这一层用于将输入的特性图进行下采样，雷同于卷积神经网络中的池化层。

假如图像的高度或宽度是奇数，PatchMerging 会进行添补，使得其变为偶数。这是因为下采样操作必要将图像分割为以2为步长的地区。假如图像的高度或宽度是奇数，直接进行切片会导致不匀称的分割，因此必要添补以保证每个块的大小一致。
这里我们在吧如上图的相同颜色块提取并进行拼接，沿着通道维度合并成一个更大的特性，将合并后的张量重新调解形状，新的空间分辨率是原来的一半（H/2 和 W/2）。

1. class PatchMerging(nn.Module):
2.     def __init__(self, dim, norm_layer=LayerNorm):
3.         super().__init__()
4.         self.dim = dim
5.         self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
6.         self.norm = norm_layer(4 * dim)
8.     def forward(self, x, H, W):
9.         """
10.         Args:
11.             x: Input feature, tensor size (B, H*W, C).
12.             H, W: Spatial resolution of the input feature.
13.         """
14.         B, L, C = x.shape
15.         assert L == H * W, "input feature has wrong size"
17.         x = x.view(B, H, W, C)
19.         if H % 2 == 1 or W % 2 == 1:
20.             x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
22.         x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
23.         x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
24.         x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
25.         x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
26.         x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
27.         x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
29.         x = self.norm(x)
30.         x = self.reduction(x)
32.         return x
34. if __name__=="__main__":
35.     batch_size = 1
36.     in_channels = 3
37.     height, width = 30, 32
38.     patch_size = 4
39.     embed_dim = 96
41.     x = torch.randn(batch_size, in_channels, height, width)
42.     patch_partition = PatchPartition(patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim)
43.     output, Wh, Ww = patch_partition(x)
45.     patch_merging = PatchMerging(dim=embed_dim)
46.     output = patch_merging(output, Wh, Ww)
47.     print(output.shape)

复制代码

在代码中呢就是在高和宽的维度通过切片的形式获得，x0表示的是左上角，x1表示的是右上角，x2表示的是左下角，x3表示的是右下角。经过一系列操作后，最后通过线性层实现通道数翻倍。
W-MSA

W-MSA（Window-based Multi-Head Self-Attention）是Swin Transformer中的一个核心创新，它是为了优化传统自注意力机制在高分辨率输入图像处理中的效率问题而提出的。

这是原论文当中给出的计算公式，h，w和C分别表示特性的高度，宽度和深度，M表示窗口的大小。在标准的 Transformer 模子中，自注意力机制必要对整个输入进行计算，这使得计算和内存的斲丧随着输入的增大而急剧增长。而在图像任务中，输入图像往往具有非常高的分辨率，因此直接应用标准的全局自注意力在计算上不可行。
W-MSA 通过在局部窗口内进行自注意力计算来解决这一问题，极大地镌汰了计算和内存开销，同时保持了模子的表示能力。
SW-MSA

SW-MSA （Shifted Window-based Multi-Head Self-Attention）结合了局部窗口化自注意力和窗口偏移（shifted）策略，既提升了计算效率，又能在捕捉局部信息的底子上，保持对全局信息的建模能力。

左侧就是刚刚说到的W-MSA，经过窗口的偏移酿成了右边的SW-MSA，偏移的策略可以或许让模子在每一层的计算中捕捉到不同窗口之间的依赖关系，制止了 W-MSA 只能在单一窗口内计算的局限。这样，相邻窗口之间的信息就可以或许通过偏移和交错的方式进行交换，增强了模子的全局感知能力。
但是，现在的窗口从原来的四个酿成了九个，假如对每一个窗口再进行W-MSA那就太麻烦了。为了应对这种情况，作者提出了一种 高效批处理计算方法，旨在优化窗口偏移后的大规模窗口计算。其核心头脑是：通过批处理计算的方式来有效地处理这些偏移后的窗口，而不是每个窗口单独计算。

意思就是说将图中的A，B，C的位置通过偏移和交错方式变化后，可以将这些窗口的计算同一进行批处理，而不是一个一个地处理。这样可以明显镌汰计算时间和内存占用。
这个过程我个人感觉比力像是卡诺图，具体的过程可以看我下面画的图：

然后这里的4还和原来的一样，5和3组合成一个窗口，1和7组合成一个窗口，8、2、6、0又组合成一个窗口，这样就和原来一样是4个4x4的窗口了，保证了计算量的不变。但是假如这样做了就会将不相邻的信息混合在一起了。作者这里接纳掩蔽机制将自注意力计算限定在每个子窗口内，实在就是创建一个蒙板来屏蔽信息。
Relative Position Bias

关于这一部分，作者没有怎么提，只是经过了相对位置偏移，指标有明显的提示。

关于这一部分，我是参考的官方代码以及b站的讲解视频理解的。首先必要创建一个相对位置偏置的参数表，它的范围是从[-Wh+1, Wh-1]，这里的 +1 和 -1 是因为偏移量是相对于当前元素的位置而言的，当前元素自身的偏移量为0，但我们不包罗0在偏移量的计算中（因为0表示没有偏移，通常会在自注意力机制中以其他方式处理）。因此，对于垂直方向（或水平方向），总的偏移量数量是 win_h（或 win_w）的正偏移量数量加上 win_h（或 win_w）的负偏移量数量，再减去一个（因为我们不计算0偏移量）。因此，相对位置偏置表的尺寸为：
   [(2 * Wh - 1) * (2 * Ww - 1), num_heads]
  每个元素的查询（Query）和键（Key）之间的内积会得到一个相似度分数，在这些分数的底子上，会加入相对位置偏置，调解相似度：
   Attention = softmax((QK^T + Relative_Position_Bias) / sqrt(d_k))
  其中，Q 是查询向量，K 是键向量，Relative_Position_Bias 是根据相对位置计算得到的偏置。加入相对位置偏置后，模子可以更好地捕捉到局部布局的依赖关系。
网络实现

1. """
2. Copyright (c) 2025, Auorui.
3. All rights reserved.
5. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
6.     <https://arxiv.org/pdf/2103.14030>
7. use for reference: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/swin_transformer.py
8.                    https://github.com/WZMIAOMIAO/deep-learning-for-image-processing/blob/master/pytorch_classification/swin_transformer/model.py
9. """
10. from functools import partial
12. import torch
13. import numpy as np
14. import torch.nn as nn
15. import torch.nn.functional as F
16. from pyzjr.utils.FormatConver import to_2tuple
17. from pyzjr.nn.models.bricks.drop import DropPath
18. from pyzjr.nn.models.bricks.initer import trunc_normal_
20. LayerNorm = partial(nn.LayerNorm, eps=1e-6)
22. class PatchPartition(nn.Module):
23.     def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None):
24.         super().__init__()
25.         self.patch_size = to_2tuple(patch_size)
26.         self.embed_dim = embed_dim
27.         self.proj = nn.Conv2d(in_channels, self.embed_dim,
28.                               kernel_size=self.patch_size, stride=self.patch_size)
29.         self.norm = norm_layer(self.embed_dim) if norm_layer else nn.Identity()
31.     def forward(self, x):
32.         _, _, H, W = x.shape
33.         if H % self.patch_size[0] != 0:
34.             pad_h = self.patch_size[0] - H % self.patch_size[0]
35.             x = F.pad(x, (0, 0, 0, pad_h))
36.         if W % self.patch_size[1] != 0:
37.             pad_w = self.patch_size[1] - W % self.patch_size[1]
38.             x = F.pad(x, (0, pad_w, 0, 0))
39.         x = self.proj(x)     # [B, embed_dim, H/patch_size, W/patch_size]
40.         Wh, Ww = x.shape[2:]
41.         x = x.flatten(2).transpose(1, 2)    # [B, num_patches, embed_dim]
42.         # Linear Embedding
43.         x = self.norm(x)
44.         # x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
45.         return x, Wh, Ww
47. class MLP(nn.Module):
48.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop_ratio=0.):
49.         super().__init__()
50.         out_features = out_features or in_features
51.         hidden_features = hidden_features or in_features
52.         self.fc1 = nn.Linear(in_features, hidden_features)
53.         self.act = act_layer()
54.         self.fc2 = nn.Linear(hidden_features, out_features)
55.         self.drop = nn.Dropout(drop_ratio)
57.     def forward(self, x):
58.         x = self.fc1(x)
59.         x = self.act(x)
60.         x = self.drop(x)
61.         x = self.fc2(x)
62.         x = self.drop(x)
63.         return x
65. class PatchMerging(nn.Module):
66.     def __init__(self, dim, norm_layer=LayerNorm):
67.         super().__init__()
68.         self.dim = dim
69.         self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
70.         self.norm = norm_layer(4 * dim)
72.     def forward(self, x, H, W):
73.         """
74.         Args:
75.             x: Input feature, tensor size (B, H*W, C).
76.             H, W: Spatial resolution of the input feature.
77.         """
78.         B, L, C = x.shape
79.         assert L == H * W, "input feature has wrong size"
81.         x = x.view(B, H, W, C)
83.         if H % 2 == 1 or W % 2 == 1:
84.             x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
86.         x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
87.         x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
88.         x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
89.         x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
90.         x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
91.         x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
93.         x = self.norm(x)
94.         x = self.reduction(x)
96.         return x
98. class WindowAttention(nn.Module):
99.     """
100.     Window based multi-head self attention (W-MSA) module with relative position bias.
101.     It supports shifted and non-shifted windows.
102.     """
103.     def __init__(
104.             self,
105.             dim,
106.             window_size,
107.             num_heads,
108.             qkv_bias=True,
109.             proj_bias=True,
110.             attention_dropout_ratio=0.,
111.             proj_drop=0.,
112.     ):
113.         super().__init__()
114.         self.dim = dim
115.         self.window_size = to_2tuple(window_size)
116.         win_h, win_w = self.window_size
117.         self.num_heads = num_heads
118.         head_dim = dim // num_heads
119.         self.scale = head_dim ** -0.5
120.         # define a parameter table of relative position bias
121.         self.relative_position_bias_table = nn.Parameter(
122.             torch.zeros((2 * win_h - 1) * (2 * win_w - 1), num_heads)
123.         )   # [2*Wh-1 * 2*Ww-1, nHeads]   Offset Range: -Wh+1, Wh-1
125.         self.register_buffer("relative_position_index",
126.                              self.get_relative_position_index(win_h, win_w), persistent=False)
127.         trunc_normal_(self.relative_position_bias_table, std=.02)
128.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
129.         self.attn_drop = nn.Dropout(attention_dropout_ratio)
130.         self.proj = nn.Linear(dim, dim, bias=proj_bias)
131.         self.proj_drop = nn.Dropout(proj_drop)
133.         self.softmax = nn.Softmax(dim=-1)
135.     def get_relative_position_index(self, win_h: int, win_w: int):
136.         # get pair-wise relative position index for each token inside the window
137.         coords = torch.stack(torch.meshgrid(torch.arange(win_h), torch.arange(win_w), indexing='ij'))  # 2, Wh, Ww
138.         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
139.         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
140.         relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
141.         relative_coords[:, :, 0] += win_h - 1  # shift to start from 0
142.         relative_coords[:, :, 1] += win_w - 1
143.         relative_coords[:, :, 0] *= 2 * win_w - 1
144.         return relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
146.     def forward(self, x, mask=None):
147.         """
148.         Args:
149.             x: input features with shape of (num_windows*B, N, C)
150.             mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
151.         """
152.         B, N, C = x.shape
153.         qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
154.         q, k, v = qkv[:3]
156.         q = q * self.scale
157.         attn = (q @ k.transpose(-2, -1))
159.         relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
160.             self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
161.         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
162.         attn = attn + relative_position_bias.unsqueeze(0)
164.         if mask is not None:
165.             nW = mask.shape[0]
166.             attn = attn.view(B // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
167.             attn = attn.view(-1, self.num_heads, N, N)
168.             attn = self.softmax(attn)
169.         else:
170.             attn = self.softmax(attn)
172.         attn = self.attn_drop(attn)
174.         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
175.         x = self.proj(x)
176.         x = self.proj_drop(x)
177.         return x
180. def window_partition(x, window_size: int):
181.     """
182.     将feature map按照window_size划分成一个个没有重叠的window
183.     Args:
184.         x: (B, H, W, C)
185.         window_size (int): window size(M)
187.     Returns:
188.         windows: (num_windows*B, window_size, window_size, C)
189.     """
190.     B, H, W, C = x.shape
191.     x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
192.     # permute: [B, H//Mh, Mh, W//Mw, Mw, C] -> [B, H//Mh, W//Mh, Mw, Mw, C]
193.     # view: [B, H//Mh, W//Mw, Mh, Mw, C] -> [B*num_windows, Mh, Mw, C]
194.     windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
195.     return windows
198. def window_reverse(windows, window_size: int, H: int, W: int):
199.     """
200.     将一个个window还原成一个feature map
201.     Args:
202.         windows: (num_windows*B, window_size, window_size, C)
203.         window_size (int): Window size(M)
204.         H (int): Height of image
205.         W (int): Width of image
207.     Returns:
208.         x: (B, H, W, C)
209.     """
210.     B = int(windows.shape[0] / (H * W / window_size / window_size))
211.     # view: [B*num_windows, Mh, Mw, C] -> [B, H//Mh, W//Mw, Mh, Mw, C]
212.     x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
213.     # permute: [B, H//Mh, W//Mw, Mh, Mw, C] -> [B, H//Mh, Mh, W//Mw, Mw, C]
214.     # view: [B, H//Mh, Mh, W//Mw, Mw, C] -> [B, H, W, C]
215.     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
216.     return x
218. class SwinTransformerBlock(nn.Module):
219.     r""" Swin Transformer Block."""
220.     mlp_ratio = 4
221.     def __init__(
222.         self,
223.         dim,
224.         num_heads,
225.         window_size=7,
226.         shift_size=0,
227.         qkv_bias=True,
228.         proj_bias=True,
229.         attention_dropout_ratio=0.,
230.         proj_drop=0.,
231.         drop_path_ratio=0.,
232.         norm_layer=LayerNorm,
233.         act_layer=nn.GELU,
234.     ):
235.         super(SwinTransformerBlock, self).__init__()
236.         self.dim = dim
237.         self.num_heads = num_heads
238.         self.window_size = window_size
239.         self.shift_size = shift_size
240.         assert 0 <= self.shift_size < window_size, "shift_size must in 0-window_size"
242.         self.norm1 = norm_layer(dim)
243.         self.attn = WindowAttention(
244.             dim,
245.             window_size=self.window_size,
246.             num_heads=num_heads,
247.             qkv_bias=qkv_bias,
248.             proj_bias=proj_bias,
249.             attention_dropout_ratio=attention_dropout_ratio,
250.             proj_drop=proj_drop,
251.         )
253.         self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
254.         self.norm2 = norm_layer(dim)
255.         mlp_hidden_dim = int(dim * self.mlp_ratio)
256.         self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop_ratio=proj_bias)
258.         self.H = None
259.         self.W = None
261.     def forward(self, x, mask_matrix):
262.         """
263.         Args:
264.             x: Input feature, tensor size (B, H*W, C).
265.             H, W: Spatial resolution of the input feature.
266.             mask_matrix: Attention mask for cyclic shift.
267.         """
268.         B, L, C = x.shape
269.         H, W = self.H, self.W
270.         assert L == H * W, "input feature has wrong size"
272.         shortcut = x
273.         x = self.norm1(x)
274.         x = x.view(B, H, W, C)
276.         # pad feature maps to multiples of window size
277.         pad_l = pad_t = 0
278.         pad_r = (self.window_size - W % self.window_size) % self.window_size
279.         pad_b = (self.window_size - H % self.window_size) % self.window_size
280.         x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
281.         _, Hp, Wp, _ = x.shape
283.         # cyclic shift
284.         if self.shift_size > 0:
285.             shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
286.             attn_mask = mask_matrix
287.         else:
288.             shifted_x = x
289.             attn_mask = None
291.         # partition windows
292.         x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
293.         x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
295.         # W-MSA/SW-MSA
296.         attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
298.         # merge windows
299.         attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
300.         shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
302.         # reverse cyclic shift
303.         if self.shift_size > 0:
304.             x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
305.         else:
306.             x = shifted_x
308.         if pad_r > 0 or pad_b > 0:
309.             x = x[:, :H, :W, :].contiguous()
311.         x = x.view(B, H * W, C)
313.         # FFN
314.         x = shortcut + self.drop_path(x)
315.         x = x + self.drop_path(self.mlp(self.norm2(x)))
317.         return x
320. class BasicLayer(nn.Module):
321.     """ A basic Swin Transformer layer for one stage."""
322.     def __init__(self,
323.                  dim,
324.                  num_layers,
325.                  num_heads,
326.                  drop_path,
327.                  window_size=7,
328.                  qkv_bias=True,
329.                  proj_bias=True,
330.                  attention_dropout_ratio=0.,
331.                  proj_drop=0.,
332.                  norm_layer=LayerNorm,
333.                  act_layer=nn.GELU,
334.                  downsample=None):
335.         super().__init__()
336.         self.window_size = window_size
337.         self.shift_size = window_size // 2
338.         self.num_layers = num_layers
340.         # build blocks
341.         self.blocks = nn.ModuleList([
342.             SwinTransformerBlock(
343.                 dim=dim,
344.                 num_heads=num_heads,
345.                 window_size=window_size,
346.                 shift_size=0 if (i % 2 == 0) else window_size // 2,
347.                 qkv_bias=qkv_bias,
348.                 proj_bias=proj_bias,
349.                 attention_dropout_ratio=attention_dropout_ratio,
350.                 proj_drop=proj_drop,
351.                 drop_path_ratio=drop_path[i] if isinstance(drop_path, list) else drop_path,
352.                 norm_layer=norm_layer,
353.                 act_layer=act_layer)
354.             for i in range(num_layers)])
355.         # patch merging layer
356.         if downsample is not None:
357.             self.downsample = downsample(dim=dim, norm_layer=norm_layer)
358.         else:
359.             self.downsample = None
361.     def forward(self, x, H, W):
362.         """ Forward function.
364.         Args:
365.             x: Input feature, tensor size (B, H*W, C).
366.             H, W: Spatial resolution of the input feature.
367.         """
369.         # calculate attention mask for SW-MSA
370.         Hp = int(np.ceil(H / self.window_size)) * self.window_size
371.         Wp = int(np.ceil(W / self.window_size)) * self.window_size
372.         img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1
373.         h_slices = (slice(0, -self.window_size),
374.                     slice(-self.window_size, -self.shift_size),
375.                     slice(-self.shift_size, None))
376.         w_slices = (slice(0, -self.window_size),
377.                     slice(-self.window_size, -self.shift_size),
378.                     slice(-self.shift_size, None))
379.         cnt = 0
380.         for h in h_slices:
381.             for w in w_slices:
382.                 img_mask[:, h, w, :] = cnt
383.                 cnt += 1
385.         mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
386.         mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
387.         attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
388.         attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
390.         for blk in self.blocks:
391.             blk.H, blk.W = H, W
392.             x = blk(x, attn_mask)
393.         if self.downsample is not None:
394.             x = self.downsample(x, H, W)
395.             H, W = (H + 1) // 2, (W + 1) // 2
397.         return x, H, W
402. class SwinTransformer(nn.Module):
403.     """ Swin Transformer backbone."""
404.     def __init__(self,
405.                  patch_size=4,
406.                  in_channels=3,
407.                  num_classes=1000,
408.                  embed_dim=96,
409.                  depths=(2, 2, 6, 2),
410.                  num_heads=(3, 6, 12, 24),
411.                  window_size=7,
412.                  qkv_bias=True,
413.                  proj_bias=True,
414.                  attention_dropout_ratio=0.,
415.                  proj_drop=0.,
416.                  drop_path_rate=0.2,
417.                  norm_layer=LayerNorm,
418.                  patch_norm=True,
419.                  ):
420.         super().__init__()
421.         self.num_classes = num_classes
422.         self.num_layers = len(depths)
423.         self.num_layers = len(depths)
424.         self.embed_dim = embed_dim
425.         self.patch_norm = patch_norm
426.         # stage4输出特征矩阵的channels
427.         self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
429.         # split image into non-overlapping patches
430.         self.patch_embed = PatchPartition(
431.             patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim,
432.             norm_layer=norm_layer if self.patch_norm else None)
434.         self.pos_drop = nn.Dropout(p=proj_drop)
435.         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
436.         layers = []
437.         for i_layer in range(self.num_layers):
438.             layer = BasicLayer(
439.                 dim=int(embed_dim * 2 ** i_layer),
440.                 num_layers=depths[i_layer],
441.                 num_heads=num_heads[i_layer],
442.                 window_size=window_size,
443.                 qkv_bias=qkv_bias,
444.                 proj_bias=proj_bias,
445.                 attention_dropout_ratio=attention_dropout_ratio,
446.                 proj_drop=proj_drop,
447.                 drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
448.                 norm_layer=norm_layer,
449.                 downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
450.                 )
451.             layers.append(layer)
453.         self.layers = nn.Sequential(*layers)
455.         self.norm = norm_layer(self.num_features)
456.         self.avgpool = nn.AdaptiveAvgPool1d(1)
457.         self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
459.         self._initialize_weights()
461.     def _initialize_weights(self):
462.         for m in self.modules():
463.             if isinstance(m, nn.Linear):
464.                 trunc_normal_(m.weight, std=.02)
465.                 if isinstance(m, nn.Linear) and m.bias is not None:
466.                     nn.init.constant_(m.bias, 0)
467.             elif isinstance(m, nn.LayerNorm):
468.                 nn.init.constant_(m.bias, 0)
469.                 nn.init.constant_(m.weight, 1.0)
471.     def forward(self, x):
472.         # x: [B, L, C]
473.         x, H, W = self.patch_embed(x)
474.         x = self.pos_drop(x)
476.         for layer in self.layers:
477.             x, H, W = layer(x, H, W)
479.         x = self.norm(x)  # [B, L, C]
480.         x = self.avgpool(x.transpose(1, 2))
481.         x = torch.flatten(x, 1)
482.         x = self.head(x)
483.         return x
485. def swin_t(num_classes) -> SwinTransformer:
486.     model = SwinTransformer(in_channels=3,
487.                             patch_size=4,
488.                             window_size=7,
489.                             embed_dim=96,
490.                             depths=(2, 2, 6, 2),
491.                             num_heads=(3, 6, 12, 24),
492.                             num_classes=num_classes)
493.     return model
495. def swin_s(num_classes) -> SwinTransformer:
496.     model = SwinTransformer(in_channels=3,
497.                             patch_size=4,
498.                             window_size=7,
499.                             embed_dim=96,
500.                             depths=(2, 2, 18, 2),
501.                             num_heads=(3, 6, 12, 24),
502.                             num_classes=num_classes)
503.     return model
506. def swin_b(num_classes) -> SwinTransformer:
507.     model = SwinTransformer(in_channels=3,
508.                             patch_size=4,
509.                             window_size=7,
510.                             embed_dim=128,
511.                             depths=(2, 2, 18, 2),
512.                             num_heads=(4, 8, 16, 32),
513.                             num_classes=num_classes)
514.     return model
516. def swin_l(num_classes) -> SwinTransformer:
517.     model = SwinTransformer(in_channels=3,
518.                             patch_size=4,
519.                             window_size=7,
520.                             embed_dim=192,
521.                             depths=(2, 2, 18, 2),
522.                             num_heads=(6, 12, 24, 48),
523.                             num_classes=num_classes)
524.     return model
526. if __name__=="__main__":
527.     import pyzjr
528.     device = 'cuda' if torch.cuda.is_available() else 'cpu'
529.     input = torch.ones(2, 3, 224, 224).to(device)
530.     net = swin_l(num_classes=4)
531.     net = net.to(device)
532.     out = net(input)
533.     print(out)
534.     print(out.shape)
535.     pyzjr.summary_1(net, input_size=(3, 224, 224))
536.     # swin_t Total params: 27,499,108
537.     # swin_s Total params: 48,792,676
538.     # swin_b Total params: 86,683,780
539.     # swin_l Total params: 194,906,308

复制代码

参考文章

Swin-Transformer网络布局详解_swin transformer-CSDN博客
Swin-transformer详解_swin transformer-CSDN博客 
【深度学习】详解 Swin Transformer (SwinT)-CSDN博客 
推荐的视频：12.1 Swin-Transformer网络布局详解_哔哩哔哩_bilibili 

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。