PyTorch 图像分割模型教程

在图像分割任务中，目标是将图像的每个像素归类为某一类，以分割出特定的物体。PyTorch 提供了非常灵活的工具，可以用于构建和训练图像分割模型。我们将使用 PyTorch 的经典网络架构，如 UNet 和 DeepLabV3，并演示如何构建、训练和测试这些模型。

1. 图像分割概述

图像分割的目标是将图像的每个像素进行分类。常见的应用场景有医学图像分割（如肿瘤检测）、自动驾驶（道路、车辆、行人分割）等。

语义分割：每个像素被分配给某个类别，例如道路、天空或车辆。
实例分割：不仅对物体分类，还要区分物体实例，如区分不同的行人。

PyTorch 中有许多预训练的模型可以直接用于图像分割任务，常用的模型包括 UNet、FCN (Fully Convolutional Network)、DeepLabV3 等。

2. 官方文档链接

PyTorch 官方文档
Torchvision 模型

3. 准备工作

在开始训练之前，我们需要安装 torch, torchvision 和 PIL 等依赖项，并准备图像数据集。您可以使用自己的图像数据集，或者使用 COCO、VOC 等常用数据集。

pip install torch torchvision pillow

4. 使用预训练的 DeepLabV3 模型

DeepLabV3 是一个性能优异的语义分割模型，PyTorch 的 torchvision 提供了预训练的 DeepLabV3 模型。我们将使用 COCO 数据集中的预训练模型，并进行推理和测试。

import torch
from torchvision import models, transforms
from PIL import Image
import matplotlib.pyplot as plt# 加载预训练的 DeepLabV3 模型
model = models.segmentation.deeplabv3_resnet50(pretrained=True)
model.eval()  # 切换到评估模式# 定义预处理步骤
preprocess = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])# 加载图像
input_image = Image.open("test_image.jpg")
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)  # 创建 batch 维度# 将输入移到 GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
input_batch = input_batch.to(device)# 进行预测
with torch.no_grad():output = model(input_batch)['out'][0]  # DeepLabV3 的输出包含 'out' 字段# 将输出转换为类别索引（每个像素对应一个类别）
output_predictions = output.argmax(0).cpu().numpy()# 显示分割结果
plt.imshow(output_predictions)
plt.show()

说明：

models.segmentation.deeplabv3_resnet50(pretrained=True)：加载使用 ResNet-50 作为主干网络的 DeepLabV3 模型，预训练于 COCO 数据集。
preprocess：对输入图像进行预处理，包括调整大小、裁剪、归一化等。
output_predictions：模型的输出是每个像素的类别索引，经过 argmax 操作，获取每个像素的类别。

5. UNet 模型

UNet 是一个广泛用于医学图像分割的经典模型。我们将从头实现 UNet 模型，并在简单的合成数据上进行训练。

5.1 UNet 网络结构

import torch
import torch.nn as nn
import torch.nn.functional as Fclass UNet(nn.Module):def __init__(self):super(UNet, self).__init__()# 下采样（编码器部分）self.encoder1 = self.double_conv(1, 64)self.encoder2 = self.double_conv(64, 128)self.encoder3 = self.double_conv(128, 256)self.encoder4 = self.double_conv(256, 512)# 中间部分self.middle = self.double_conv(512, 1024)# 上采样（解码器部分）self.upconv4 = self.up_conv(1024, 512)self.decoder4 = self.double_conv(1024, 512)self.upconv3 = self.up_conv(512, 256)self.decoder3 = self.double_conv(512, 256)self.upconv2 = self.up_conv(256, 128)self.decoder2 = self.double_conv(256, 128)self.upconv1 = self.up_conv(128, 64)self.decoder1 = self.double_conv(128, 64)# 最后的分类层self.final = nn.Conv2d(64, 1, kernel_size=1)def double_conv(self, in_channels, out_channels):return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),nn.ReLU(inplace=True),nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),nn.ReLU(inplace=True),)def up_conv(self, in_channels, out_channels):return nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)def forward(self, x):# 编码器部分e1 = self.encoder1(x)e2 = self.encoder2(F.max_pool2d(e1, 2))e3 = self.encoder3(F.max_pool2d(e2, 2))e4 = self.encoder4(F.max_pool2d(e3, 2))# 中间部分middle = self.middle(F.max_pool2d(e4, 2))# 解码器部分d4 = self.upconv4(middle)d4 = torch.cat((e4, d4), dim=1)d4 = self.decoder4(d4)d3 = self.upconv3(d4)d3 = torch.cat((e3, d3), dim=1)d3 = self.decoder3(d3)d2 = self.upconv2(d3)d2 = torch.cat((e2, d2), dim=1)d2 = self.decoder2(d2)d1 = self.upconv1(d2)d1 = torch.cat((e1, d1), dim=1)d1 = self.decoder1(d1)return self.final(d1)# 创建模型实例
unet_model = UNet()
print(unet_model)

说明：

UNet 是一种编码-解码结构，包含多个下采样（编码器）和上采样（解码器）层。每次下采样都会减少特征图的大小，并增加特征通道数，上采样则恢复原始图像的大小。
ConvTranspose2d 用于进行上采样操作。

5.2 训练 UNet 模型

为了训练 UNet 模型，我们需要构建一个数据加载器并定义损失函数和优化器。我们以一个简单的二分类分割任务为例。

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms# 创建合成数据集
class SyntheticSegmentationDataset(Dataset):def __init__(self, num_samples, image_size):self.num_samples = num_samplesself.image_size = image_sizeself.transform = transforms.Compose([transforms.ToTensor(),])def __len__(self):return self.num_samplesdef __getitem__(self, idx):image = torch.rand(1, self.image_size, self.image_size)mask = (image > 0.5).float()  # 简单的二分类掩码return image, mask# 创建数据集
dataset = SyntheticSegmentationDataset(num_samples=1000, image_size=128)
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)# 定义损失函数和优化器
criterion = nn.BCEWithLogitsLoss()  # 二分类交叉熵损失
optimizer = torch.optim.Adam(unet_model.parameters(), lr=0.001)# 训练循环
unet_model.train()
for epoch in range(5):  # 简单训练 5 个 epochfor images, masks in dataloader:# 前向传播outputs = unet_model(images)loss = criterion(outputs, masks)# 反向传播和优化optimizer.zero_grad()loss.backward()optimizer.step()print(f'Epoch [{epoch+1}/5], Loss: {loss.item():.4f}')