swin transformer pytorch代码实现

Swin Transformer 是一种新型的 Transformer 模型，它在计算效率和模型性能之间取得了很好的平衡。以下是使用 PyTorch 实现 Swin Transformer 的代码示例：

import torch
import torch.nn as nn
from einops.layers.torch import Rearrange

class SwinBlock(nn.Module):
    def __init__(self, in_channels, out_channels, window_size=7, shift_size=2):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.norm1 = nn.BatchNorm2d(out_channels)
        self.window_size = window_size
        self.shift_size = shift_size
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=window_size, stride=shift_size, padding=window_size//2, groups=out_channels)
        self.norm2 = nn.BatchNorm2d(out_channels)
        self.conv3 = nn.Conv2d(out_channels, out_channels, kernel_size=1)
        self.norm3 = nn.BatchNorm2d(out_channels)
        self.downsample = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=1),
            nn.BatchNorm2d(out_channels)
        )
        
    def forward(self, x):
        identity = x
        x = self.conv1(x)
        x = self.norm1(x)
        x = nn.functional.relu(x)
        x = Rearrange(x, 'b c h w -> b (h w) c')
        x = Rearrange(x, 'b (h w) c -> b h w c', h=int(x.shape[1]**0.5))
        x = self.conv2(x)
        x = self.norm2(x)
        x = nn.functional.relu(x)
        x = self.conv3(x)
        x = self.norm3(x)
        if identity.shape != x.shape:
            identity = self.downsample(identity)
        x += identity
        return x

class SwinTransformer(nn.Module):
    def __init__(self, in_channels=3, num_classes=1000, hidden_dim=96, layers=[2, 2, 6, 2], window_sizes=[7, 3, 3, 3], shift_sizes=[4, 2, 2, 2]):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(in_channels, hidden_dim, kernel_size=4, stride=4),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.layers = nn.ModuleList()
        for i in range(len(layers)):
            layer = nn.ModuleList()
            for j in range(layers[i]):
                if j == 0:
                    block = SwinBlock(hidden_dim*(2**i), hidden_dim*(2**i), window_sizes[i], shift_sizes[i])
                else:
                    block = SwinBlock(hidden_dim*(2**i), hidden_dim*(2**i+1), window_sizes[i], shift_sizes[i])
                layer.append(block)
            self.layers.append(layer)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(hidden_dim*(2**len(layers)), num_classes)
        
    def forward(self, x):
        x = self.stem(x)
        for i in range(len(self.layers)):
            for j in range(len(self.layers[i])):
                x = self.layers[i][j](x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

这是一个简单的 Swin Transformer 实现，包含了 SwinBlock 和 SwinTransformer 两个类。其中 SwinBlock 是 Swin Transformer 的基本模块，SwinTransformer 是整个模型的实现。你可以根据自己的需求进行修改和扩展。