UNet网络详解及PyTorch实现

一、UNet网络原理

U-Net，自2015年诞生以来，便以其卓越的性能在生物医学图像分割领域崭露头角。作为FCN的一种变体，U-Net凭借其Encoder-Decoder的精巧结构，不仅在医学图像分析中大放异彩，更在卫星图像分割、工业瑕疵检测等多个领域展现出强大的应用能力。UNet是一种常用于图像分割的卷积神经网络架构，其特点在于其U型结构，包括一个收缩路径（下采样）和一个扩展路径（上采样）。这种结构使得UNet能够在捕获上下文信息的同时，也能精确地定位到目标边界。

收缩路径（编码器Encoder）：通过连续的卷积和池化操作，逐步减小特征图的尺寸，从而捕获到图像的上下文信息。

扩展路径（解码器Decoder）：通过上采样操作逐步恢复特征图的尺寸，并与收缩路径中对应尺度的特征图进行拼接（concatenate），以融合不同尺度的特征信息。

跳跃连接：UNet中的跳跃连接使得扩展路径能够利用到收缩路径中的高分辨率特征，从而提高了分割的精度。

输出层：UNet的输出层通常是一个1x1的卷积层，用于将特征图转换为与输入图像相同尺寸的分割图。

二、基于PyTorch的UNet实现

下面是一个简单的基于PyTorch的UNet实现，用于图像分割任务。(环境安装可以看我往期博客)

<code>import torch

import torch.nn as nn

import torch.nn.functional as F

class DoubleConv(nn.Module):

"""(convolution => [BN] => ReLU) * 2"""

def __init__(self, in_channels, out_channels):

super(DoubleConv, self).__init__()

self.conv = nn.Sequential(

nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),

nn.BatchNorm2d(out_channels),

nn.ReLU(inplace=True),

nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),

nn.BatchNorm2d(out_channels),

nn.ReLU(inplace=True)

)

def forward(self, x):

return self.conv(x)

class UNet(nn.Module):

def __init__(self, n_channels, n_classes, bilinear=True):

super(UNet, self).__init__()

self.n_channels = n_channels

self.n_classes = n_classes

self.bilinear = bilinear

self.inc = DoubleConv(n_channels, 64)

self.down1 = DoubleConv(64, 128)

self.down2 = DoubleConv(128, 256)

self.down3 = DoubleConv(256, 512)

factor = 2 if bilinear else 1

self.down4 = DoubleConv(512, 1024 // factor)

self.up1 = nn.ConvTranspose2d(1024 // factor, 512 // factor, kernel_size=2, stride=2)

self.up2 = nn.ConvTranspose2d(512 // factor, 256 // factor, kernel_size=2, stride=2)

self.up3 = nn.ConvTranspose2d(256 // factor, 128 // factor, kernel_size=2, stride=2)

self.up4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)

self.outc = nn.Conv2d(64, n_classes, kernel_size=1)

def forward(self, x):

x1 = self.inc(x)

x2 = self.down1(x1)

x3 = self.down2(x2)

x4 = self.down3(x3)

x5 = self.down4(x4)

x = self.up1(x5, x4)

x = self.up2(x, x3)

x = self.up3(x, x2)

x = self.up4(x, x1)

logits = self.outc(x)

return logits

三、训练与推理的完整代码

首先，我们需要准备数据集、定义损失函数和优化器，然后编写训练循环。

python

import torch

import torch.nn as nn

import torch.optim as optim

from torch.utils.data import DataLoader

from torchvision import transforms, datasets

from unet_model import UNet # 假设UNet定义在unet_model.py文件中

# 设定超参数

num_epochs = 10

learning_rate = 0.001

batch_size = 4

# 数据预处理

transform = transforms.Compose([

transforms.Resize((256, 256)),

transforms.ToTensor(),

transforms.Normalize(mean=[0.5], std=[0.5])

])

# 加载训练集

train_dataset = datasets.ImageFolder(root='path_to_train_dataset', transform=transform) code>

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 定义模型、损失函数和优化器

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

num_classes = len(train_dataset.classes) # 根据数据集确定类别数

model = UNet(n_channels=3, n_classes=num_classes).to(device)

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练循环

for epoch in range(num_epochs):

model.train() # 设置模型为训练模式

running_loss = 0.0

for i, data in enumerate(train_loader):

inputs, labels = data[0].to(device), data[1].to(device)

optimizer.zero_grad() # 梯度清零

outputs = model(inputs) # 前向传播

loss = criterion(outputs, labels) # 计算损失

loss.backward() # 反向传播

optimizer.step() # 更新权重

running_loss += loss.item() * inputs.size(0)

epoch_loss = running_loss / len(train_loader.dataset)

print(f'Epoch { epoch+1}/{ num_epochs}, Loss: { epoch_loss:.4f}')

# 保存模型

torch.save(model.state_dict(), 'unet_model.pth')

推理

在推理阶段，我们加载已训练好的模型，并对测试集或单个图像进行预测。

python

# 加载模型

model.load_state_dict(torch.load('unet_model.pth'))

model.eval() # 设置模型为评估模式

# 如果需要，准备测试集

test_dataset = datasets.ImageFolder(root='path_to_test_dataset', transform=transform) code>

test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 对测试集进行推理

with torch.no_grad():

for inputs, _ in test_loader:

inputs = inputs.to(device)

outputs = model(inputs)

_, predicted = torch.max(outputs, 1)

# 可以将predicted保存为文件或进行其他处理

# 对单个图像进行推理

image_path = 'path_to_single_image.png'

image = Image.open(image_path).convert('RGB') # 确保是RGB格式

image = transform(image).unsqueeze(0).to(device) # 对图像进行预处理并添加到batch维度

with torch.no_grad():

prediction = model(image)

_, predicted = torch.max(prediction, 1)

predicted_class = train_dataset.classes[predicted.item()] # 获取预测的类别名

# 可以将predicted保存为文件或进行可视化

这里我假设你已经有了适当的训练和测试数据集，并且它们已经被组织成了ImageFolder可以理解的格式（即每个类别的图像都在一个单独的子文件夹中）。此外，代码中的transform可能需要根据你的具体数据集进行调整。

在推理阶段，我们使用torch.max来找出每个图像最有可能的类别，并通过predicted_class变量打印或返回该类别。对于测试集，你可能希望将预测结果保存为文件，以便后续分析或可视化。对于单个图像，你可以直接进行可视化或将其保存为带有分割结果的图像。

四、总结