swin transformer+FPN（内含代码，可用于图像分类）

以下是一个基础版本的 Swin Transformer（Swin-B）加上特征金字塔网络（FPN）实现渐进融合的简化代码。请注意，这是一个简化版本，可能需要根据具体需求进行调整和优化。

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

# Swin Transformer Block
class SwinTransformerBlock(nn.Module):
    def __init__(self, dim, heads, head_dim, mlp_dim, dropout=0.0):
        super(SwinTransformerBlock, self).__init__()
        self.attention = nn.MultiheadAttention(dim, heads, dropout=dropout)
        self.norm1 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Linear(mlp_dim, dim),
            nn.Dropout(dropout)
        )
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x):
        attention_output, _ = self.attention(x, x, x)
        x = x + attention_output
        x = self.norm1(x)
        mlp_output = self.mlp(x)
        x = x + mlp_output
        x = self.norm2(x)
        return x

# Swin Transformer Backbone
class SwinTransformer(nn.Module):
    def __init__(self, image_size, patch_size, in_channels, num_classes, embed_dim, depths, heads, mlp_dim, dropout=0.0):
        super(SwinTransformer, self).__init__()
        self.patch_embed = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.pos_embed = nn.Parameter(torch.zeros(1, (image_size // patch_size) ** 2 + 1, embed_dim))
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(dim=embed_dim, heads=heads, head_dim=embed_dim // heads, mlp_dim=mlp_dim, dropout=dropout)
            for _ in range(depths)
        ])
        self.norm = nn.LayerNorm(embed_dim)
        self.fc = nn.Linear(embed_dim, num_classes)

    def forward(self, x):
        x = self.patch_embed(x)
        B, C, H, W = x.shape
        x = x.flatten(2).transpose(1, 2)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embed
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)
        cls_tokens = x[:, 0]
        output = self.fc(cls_tokens)
        return output

# Feature Pyramid Network Block
class FPNBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(FPNBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.up_sample = nn.Upsample(scale_factor=2, mode='nearest')

    def forward(self, x):
        return self.conv(self.up_sample(x))

# Swin Transformer with Feature Pyramid Network (FPN)
class SwinTransformerWithFPN(nn.Module):
    def __init__(self, image_size, patch_size, in_channels, num_classes, embed_dim, depths, heads, mlp_dim, fpn_channels, dropout=0.0):
        super(SwinTransformerWithFPN, self).__init__()

        # Swin Transformer
        self.swin = SwinTransformer(image_size, patch_size, in_channels, num_classes, embed_dim, depths, heads, mlp_dim, dropout)

        # FPN Blocks
        self.fpn_block1 = FPNBlock(embed_dim, fpn_channels)
        self.fpn_block2 = FPNBlock(embed_dim, fpn_channels)
        self.fpn_block3 = FPNBlock(embed_dim, fpn_channels)

        # Classifier
        self.classifier = nn.Linear(fpn_channels, num_classes)

    def forward(self, x):
        # Swin Transformer
        features = self.swin(x)

        # FPN
        fpn_feature1 = self.fpn_block1(features.blocks[-3])
        fpn_feature2 = self.fpn_block2(features.blocks[-4])
        fpn_feature3 = self.fpn_block3(features.blocks[-5])

        # Combine FPN features
        fused_feature = fpn_feature1 + fpn_feature2 + fpn_feature3

        # Global Average Pooling
        global_pooling = torch.mean(fused_feature, dim=[2, 3])

        # Classifier
        output = self.classifier(global_pooling)

        return output

# 创建 Swin Transformer + FPN 模型
swin_fpn_model = SwinTransformerWithFPN(
    image_size=224,
    patch_size=4,
    in_channels=3,
    num_classes=1000,
    embed_dim=96,
    depths=12,
    heads=4,
    mlp_dim=384,
    fpn_channels=64,
    dropout=0.0
)

# 打印模型结构
print(swin_fpn_model)

这个代码示例定义了一个简化版本的 Swin Transformer 和 FPN 结合的模型。你可以根据需要调整 Swin Transformer 和 FPNBlock 的通道数以适应你的任务。确保你的输入图像尺寸和通道数与模型定义中的一致。