swin_transformer代码

Swin Transformer是2021年提出的一种新型Transformer模型，其在图像分类、物体检测等任务上表现优异。以下是Swin Transformer的代码实现：

import torch
import torch.nn as nn
import torch.nn.functional as F

class SwinBlock(nn.Module):
    def __init__(self, dim, num_heads, window_size, shift_size):
        super(SwinBlock, self).__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
        self.window_size = window_size
        self.shift_size = shift_size

    def forward(self, x):
        # Shift windows and flatten them
        n, c, h, w = x.shape
        unfold = nn.Unfold(kernel_size=(self.window_size, self.window_size), 
                           stride=(self.shift_size, self.shift_size))
        windows = unfold(x).view(n, c, -1, self.window_size * self.window_size).transpose(1, 2)

        # Attention
        residual1 = windows
        windows = self.norm1(windows)
        windows, _ = self.attn(windows, windows, windows)
        windows = residual1 + windows

        # MLP
        residual2 = windows
        windows = self.norm2(windows)
        windows = self.mlp(windows)
        windows = residual2 + windows
        
        # Reshape and reassemble to original shape
        fold = nn.Fold(output_size=(h, w), kernel_size=(self.window_size, self.window_size), 
                       stride=(self.shift_size, self.shift_size))
        x = fold(windows.transpose(1, 2).contiguous().view(n, -1, c))
        return x

class SwinTransformer(nn.Module):
    def __init__(self, in_channels, num_classes, hidden_dim=96, num_blocks=2, num_heads=3, window_sizes=[7, 3], shift_sizes=[4, 2]):
        super(SwinTransformer, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=4, stride=4)
        self.norm1 = nn.LayerNorm([hidden_dim, 56, 56])
        self.blocks = nn.ModuleList([SwinBlock(hidden_dim, num_heads, window_sizes[i], shift_sizes[i]) for i in range(len(window_sizes))])
        self.norm2 = nn.LayerNorm([hidden_dim, 56, 56])
        self.conv2 = nn.Conv2d(hidden_dim, in_channels, kernel_size=1)
        self.norm3 = nn.LayerNorm([in_channels, 56, 56])
        self.fc = nn.Linear(in_channels * 56 * 56, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.norm1(x)

        for block in self.blocks:
            x = block(x)

        x = self.norm2(x)
        x = self.conv2(x)
        x = self.norm3(x)
        x = x.flatten(start_dim=1)
        x = self.fc(x)
        return x

以上代码实现了Swin Transformer的基本结构，包括SwinBlock和SwinTransformer模块。其中SwinBlock是Swin Transformer的基本单元，而SwinTransformer则是由多个SwinBlock堆叠而成的完整模型。需要注意的是，SwinTransformer的输入需要是一个四维的张量，即(batch_size, in_channels, height, width)。