用pytorch搭建Swin Transformer

Swin Transformer是一种高效的Transformer架构，适用于图像分类、目标检测等计算机视觉任务。在本文中，我们将使用PyTorch实现Swin Transformer。

首先，我们需要安装PyTorch和其他必要的库：

pip install torch torchvision
pip install einops
pip install timm

接下来，我们将定义Swin Transformer模型。我们将使用timm库中提供的预训练模型作为基础模型，并在其上添加我们自己的头部。我们将使用Swin-L模型，其中L表示层数。这里我们定义一个函数来实现这个过程：

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

def build_swin(num_classes):
    model = timm.create_model('swin_large_patch4_window12_384', pretrained=True)
    model.head = nn.Sequential(
        nn.AdaptiveAvgPool2d((1, 1)),
        Rearrange('b c h w -> b (c h w)'),
        nn.Linear(1024, num_classes)
    )
    return model

在这个函数中，我们首先使用`timm.create_model`加载预训练的Swin-L模型。然后，我们替换模型的头部，使用AdaptiveAvgPool2d对特征图进行平均池化，然后使用Rearrange层将特征图重排列为(batch_size, features)的形状，最后使用一个全连接层将其映射到类别数量。

接下来，我们将定义训练和测试函数。在训练函数中，我们将使用交叉熵损失和随机梯度下降优化器进行训练。在测试函数中，我们将计算模型在测试集上的准确率。

import torch.optim as optim

def train(model, train_loader, test_loader, criterion, optimizer, num_epochs=10):
    for epoch in range(num_epochs):
        model.train()
        for i, (images, labels) in enumerate(train_loader):
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, labels in test_loader:
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        accuracy = 100 * correct / total
        print('Epoch [{}/{}], Accuracy: {:.2f}%'
              .format(epoch + 1, num_epochs, accuracy))

def test(model, test_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    print('Accuracy: {:.2f}%'.format(accuracy))

最后，我们将加载数据集并开始训练和测试模型：

import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader

# Define transforms
transform = transforms.Compose([
    transforms.Resize((384, 384)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# Load data
train_dataset = datasets.ImageFolder('train', transform=transform)
test_dataset = datasets.ImageFolder('test', transform=transform)

# Define data loaders
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# Build and train model
model = build_swin(num_classes=2)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
train(model, train_loader, test_loader, criterion, optimizer, num_epochs=10)

# Test model
test(model, test_loader)

在这个例子中，我们使用了一个二分类任务，我们将数据集分为train和test两个文件夹中，每个文件夹包含两个文件夹，分别是类别1和类别2的图像。我们使用了384x384的输入大小，并训练了10个epoch。在测试集上，我们获得了97.50%的准确率。