segformer语义分割代码

以下是一个基于Segformer模型的简单语义分割代码，使用的是PyTorch框架。

首先，需要导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models.segmentation as segmentation

然后，定义模型：

class Segformer(nn.Module):
    def __init__(self, num_classes):
        super(Segformer, self).__init__()
        
        # 使用预训练的Segformer模型
        self.segformer = segmentation.segformer.segformer_resnet50d_256(pretrained=True, num_classes=num_classes)
        
    def forward(self, x):
        return self.segformer(x)['out']

接下来，定义训练函数：

def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

定义测试函数：

def test(model, device, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            
    test_loss /= len(test_loader.dataset)
    print('Test set: Average loss: {:.4f}'.format(test_loss))

最后，定义主函数并开始训练：

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 加载数据集
    train_dataset = datasets.Cityscapes('path/to/dataset', split='train', mode='fine', target_type='semantic')
    test_dataset = datasets.Cityscapes('path/to/dataset', split='val', mode='fine', target_type='semantic')
    
    # 数据预处理
    transform = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])
    
    train_dataset.transform = transform
    test_dataset.transform = transform
    
    # 加载数据
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
    
    # 定义模型、损失函数和优化器
    model = Segformer(num_classes=34).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())
    
    # 训练模型
    for epoch in range(1, 11):
        train(model, device, train_loader, optimizer, criterion, epoch)
        test(model, device, test_loader, criterion)

这是一个基本的Segformer语义分割代码示例，可以根据需要进行修改。