UNet 改进(5)：结合SE模块提升图像分割性能

U-Net是医学图像分割领域最成功的架构之一，其对称的编码器-解码器结构和跳跃连接使其能够有效捕捉多尺度特征。本文将解析一个改进版的U-Net实现，该版本通过引入Squeeze-and-Excitation(SE)模块进一步提升了模型性能。

一、架构概览

这个改进的U-Net保持了经典U-Net的核心结构，但在每个卷积块后添加了SE模块，主要包含以下几个关键组件：

SE注意力模块：增强重要通道的特征响应
双卷积块：基础特征提取单元
编码器-解码器结构：逐步下采样和上采样
跳跃连接：结合低层和高层特征

二、核心组件详解

1. SE注意力模块 (SELayer)

class SELayer(nn.Module):def __init__(self, in_channels, reduction=16):super(SELayer, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.fc = nn.Sequential(nn.Linear(in_channels, in_channels // reduction, bias=False),nn.ReLU(inplace=True),nn.Linear(in_channels // reduction, in_channels, bias=False),nn.Sigmoid())

SE模块通过以下步骤工作：

使用全局平均池化将空间信息压缩为一个通道描述符
通过两个全连接层学习通道间的依赖关系
使用Sigmoid激活生成通道权重
将权重应用于原始特征图

这种机制让模型能够自适应地强调重要特征通道，抑制不重要的通道。

2. 改进的双卷积块 (DoubleConv)

class DoubleConv(nn.Module):def __init__(self, in_channels, out_channels):super().__init__()self.double_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),SELayer(out_channels)  # 添加 SE 模块)

每个双卷积块包含：

两个3×3卷积层，保持空间分辨率(padding=1)
每个卷积后接批量归一化和ReLU激活
最后添加SE模块进行通道注意力加权

3. 完整的改进U-Net (ImprovedUNet)

编码器部分通过最大池化逐步下采样，解码器部分通过转置卷积上采样，并结合跳跃连接：

class ImprovedUNet(nn.Module):def __init__(self, n_channels, n_classes):# 初始化各层...def forward(self, x):# 编码过程x1 = self.inc(x)       # 初始卷积x2 = self.down1(x1)    # 下采样1x3 = self.down2(x2)    # 下采样2x4 = self.down3(x3)    # 下采样3x5 = self.down4(x4)    # 下采样4# 解码过程x = self.up1(x5)x = self.double_conv_up1(torch.cat([x, F.interpolate(x4, size=x.shape[2:])], dim=1))# ...类似处理其他上采样层return self.outc(x)

三、创新点与优势

SE模块集成：在每个双卷积块后添加SE模块，使模型能够自适应地重新校准通道特征响应
改进的特征融合：使用双线性插值调整跳跃连接特征图尺寸，确保精确对齐
参数效率：通过factor参数控制解码器通道数，平衡模型容量和计算成本

四、性能分析

这个改进版U-Net相比原始U-Net有以下潜在优势：

更好的特征选择能力，通过SE模块突出重要特征
更稳定的训练，得益于批量归一化的广泛使用
更精确的边界预测，得益于改进的特征融合方式

五、使用示例

# 创建模型实例
model = ImprovedUNet(n_channels=3, n_classes=1)# 随机输入测试
input_tensor = torch.randn(2, 3, 256, 256)  # 2张256x256的RGB图像
output = model(input_tensor)  # 输出形状为[2, 1, 256, 256]

六、完整代码

import torch
import torch.nn as nn
import torch.nn.functional as F# SE 模块
class SELayer(nn.Module):def __init__(self, in_channels, reduction=16):super(SELayer, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)self.fc = nn.Sequential(nn.Linear(in_channels, in_channels // reduction, bias=False),nn.ReLU(inplace=True),nn.Linear(in_channels // reduction, in_channels, bias=False),nn.Sigmoid())def forward(self, x):b, c, _, _ = x.size()y = self.avg_pool(x).view(b, c)y = self.fc(y).view(b, c, 1, 1)return x * y.expand_as(x)# 改进的卷积块
class DoubleConv(nn.Module):def __init__(self, in_channels, out_channels):super().__init__()self.double_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),SELayer(out_channels)  # 添加 SE 模块)def forward(self, x):return self.double_conv(x)# 改进的 U-Net 模型
class ImprovedUNet(nn.Module):def __init__(self, n_channels, n_classes):super().__init__()self.n_channels = n_channelsself.n_classes = n_classesself.inc = DoubleConv(n_channels, 64)self.down1 = nn.Sequential(nn.MaxPool2d(2),DoubleConv(64, 128))self.down2 = nn.Sequential(nn.MaxPool2d(2),DoubleConv(128, 256))self.down3 = nn.Sequential(nn.MaxPool2d(2),DoubleConv(256, 512))factor = 2self.down4 = nn.Sequential(nn.MaxPool2d(2),DoubleConv(512, 1024 // factor))self.up1 = nn.Sequential(nn.ConvTranspose2d(1024 // factor, 512 // factor, kernel_size=2, stride=2))self.double_conv_up1 = DoubleConv(512 // factor + 512, 512 // factor)self.up2 = nn.Sequential(nn.ConvTranspose2d(512 // factor, 256 // factor, kernel_size=2, stride=2))self.double_conv_up2 = DoubleConv(256 // factor + 256, 256 // factor)self.up3 = nn.Sequential(nn.ConvTranspose2d(256 // factor, 128 // factor, kernel_size=2, stride=2))self.double_conv_up3 = DoubleConv(128 // factor + 128, 128 // factor)self.up4 = nn.Sequential(nn.ConvTranspose2d(128 // factor, 64, kernel_size=2, stride=2))self.double_conv_up4 = DoubleConv(64 + 64, 64)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)x = self.double_conv_up1(torch.cat([x, F.interpolate(x4, size=x.shape[2:], mode='bilinear', align_corners=True)], dim=1))x = self.up2(x)x = self.double_conv_up2(torch.cat([x, F.interpolate(x3, size=x.shape[2:], mode='bilinear', align_corners=True)], dim=1))x = self.up3(x)x = self.double_conv_up3(torch.cat([x, F.interpolate(x2, size=x.shape[2:], mode='bilinear', align_corners=True)], dim=1))x = self.up4(x)x = self.double_conv_up4(torch.cat([x, F.interpolate(x1, size=x.shape[2:], mode='bilinear', align_corners=True)], dim=1))logits = self.outc(x)return logits# 创建改进的 U-Net 模型实例
model = ImprovedUNet(n_channels=3, n_classes=1)
print(model)# 生成一个随机输入
input_tensor = torch.randn(2, 3, 256, 256)# 前向传播
output = model(input_tensor)
print(output.shape)