Pytorch图像分割

PyTorch是一个流行的深度学习框架，可以用于图像分割任务。下面是一个简单的PyTorch图像分割的例子：

1. 导入必要的库和数据集

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms, datasets

# 加载数据集
train_data = datasets.Cityscapes('./data', split='train', mode='fine', target_type='semantic',
                                 transform=transforms.Compose([
                                     transforms.ToTensor(),
                                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                                 ]))

test_data = datasets.Cityscapes('./data', split='test', mode='fine', target_type='semantic',
                                transform=transforms.Compose([
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
                                ]))

# 加载数据集的dataloader
train_loader = DataLoader(train_data, batch_size=4, shuffle=True)
test_loader = DataLoader(test_data, batch_size=4, shuffle=False)

2. 定义模型

class Net(torch.nn.Module):
    def __init__(self, num_classes=19):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = torch.nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = torch.nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.conv5 = torch.nn.Conv2d(512, 512, kernel_size=3, padding=1)

        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv6 = torch.nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv7 = torch.nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.conv8 = torch.nn.Conv2d(512, num_classes, kernel_size=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = self.pool(F.relu(self.conv4(x)))
        x = self.pool(F.relu(self.conv5(x)))

        x = self.pool(F.relu(self.conv6(x)))
        x = self.pool(F.relu(self.conv7(x)))
        x = self.conv8(x)

        return x

3. 训练模型

# 定义模型和优化器
model = Net()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 定义损失函数
criterion = torch.nn.CrossEntropyLoss()

# 训练模型
for epoch in range(10):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets.squeeze(1).long())
        loss.backward()
        optimizer.step()

        if i % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 10, i+1, len(train_loader), loss.item()))

4. 测试模型

# 测试模型
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for inputs, targets in test_loader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += (predicted == targets.squeeze(1)).sum().item()

    print('Accuracy of the network on the test images: {:.2f} %'.format(100 * correct / total))

这是一个简单的图像分割的例子，你可以根据自己的需求进行修改和扩展。