基于Python深度学习框架Pytorch的图像分割代码示例

以下是一个基于Pytorch的图像分割代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import Cityscapes
from model import UNet

# 设置超参数
num_epochs = 10
batch_size = 4
learning_rate = 0.001

# 加载Cityscapes数据集
train_dataset = Cityscapes(root='./data', split='train', mode='fine',
                           target_type=['semantic'], transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 定义UNet模型和损失函数
model = UNet(n_channels=3, n_classes=19)
criterion = nn.CrossEntropyLoss()

# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 每100个iteration打印一次损失
        if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Iteration [{}/{}], Loss: {:.4f}'
                  .format(epoch+1, num_epochs, i+1, len(train_loader), loss.item()))

# 保存模型
torch.save(model.state_dict(), 'unet.ckpt')

这里使用了一个名为`UNet`的卷积神经网络模型进行图像分割，该模型的实现可以参考这个[代码库](https://github.com/milesial/Pytorch-UNet)。在训练过程中，我们使用了`Cityscapes`数据集，它是一个用于城市场景图像分割的数据集。在每个iteration结束后，我们打印了损失值以便于观察训练过程。最后，我们将训练好的模型保存到了`unet.ckpt`文件中。