UNet

从零开始理解UNet：模型结构与训练细节

UNet是一种经典的深度学习模型，由Olaf Ronneberger等人于2015年提出，最初用于生物医学图像分割任务。其对称的U形结构、跳跃连接设计以及对小样本数据的适应性，使其成为图像分割领域的基准模型之一。以下从模型结构和训练细节两方面展开介绍：

一、模型结构

UNet的架构分为编码器（收缩路径）、解码器（扩展路径）和跳跃连接三部分，整体呈对称的U形结构。

1. 编码器（Encoder）
   编码器通过下采样逐步提取图像的高层语义信息，同时缩小特征图的空间尺寸：

   • 卷积块：每个编码阶段包含两次3×3卷积操作，每次卷积后接批量归一化（BatchNorm）和ReLU激活函数。
   • 池化层：通过2×2最大池化（MaxPooling）将特征图尺寸减半，通道数翻倍（例如从64通道增至128通道），以扩大感受野。
   • 层级示例：输入图像（如572×572）经过4次下采样后，特征图尺寸最小可达约32×32，通道数增至1024。

2. 解码器（Decoder）
   解码器通过上采样恢复空间分辨率，并融合编码器的多尺度特征：

   • 上采样操作：使用2×2反卷积（Transposed Convolution）或双线性插值（Upsampling）扩大特征图尺寸。
   • 跳跃连接（Skip Connection）：将编码器对应层的高分辨率特征与解码器当前层的特征拼接（Concatenate），补充细节信息（如边界、纹理）。
   • 卷积块：拼接后的特征再经过两次3×3卷积，进一步融合高低层信息。

3. 输出层
   最后一层通过1×1卷积将通道数映射为类别数（如二分类输出1通道），并使用Sigmoid或Softmax激活函数生成概率图。

结构示意图

输入图像 → 编码器（下采样） → 瓶颈层 → 解码器（上采样+跳跃连接） → 分割结果

二、训练细节

UNet的训练需关注数据增强、损失函数设计、优化策略等关键点：

1. 数据准备与增强

   • 医学数据特点：样本量少、标注成本高，需依赖数据增强（如弹性变形、旋转、翻转、缩放）提升泛化能力。
   • 输入输出对齐：原始论文中输入572×572，输出388×388，需通过裁剪或填充保证尺寸匹配。

2. 损失函数设计

   • 交叉熵损失（CE Loss）：常规分类损失，但对类别不平衡敏感（如医学图像中前景占比低）。
   • Dice Loss：基于重叠区域度量（Dice系数），适用于小目标分割，公式为：
     $$ \text{Dice Loss} = 1 - \frac{2 \sum y_{\text{pred}} y_{\text{true}}}{\sum y_{\text{pred}} + \sum y_{\text{true}}} $$ 可结合交叉熵使用以兼顾全局和局部优化。
   • 边界权重：在细胞分割任务中，对相邻细胞边界区域赋予更高权重，避免误分割。

3. 优化策略

   • 优化器选择：常用Adam或带动量的SGD，初始学习率设为1e-3至1e-4，配合学习率衰减。
   • 早停法（Early Stopping）：通过验证集损失监控过拟合风险，提前终止训练。
   • 正则化：使用Dropout或权重衰减（L2正则化）抑制过拟合。

4. 评估指标

   • IoU（交并比）：衡量预测区域与真实标签的重合度。
   • Dice系数：与Dice Loss对应的评估指标，值越高表示分割精度越好。

三、UNet的改进与扩展

• UNet++：引入密集跳跃连接，增强多尺度特征融合能力。
• Attention UNet：在跳跃连接中嵌入注意力机制，聚焦关键区域。
• 3D UNet：扩展至三维体积数据（如CT、MRI），支持体素级分割。

总结

UNet通过编码器-解码器结构和跳跃连接，实现了从粗粒度到细粒度特征的多层次融合，在医学影像分割等领域表现出色。其训练需注重数据增强、损失函数设计及正则化策略。后续改进模型（如UNet++、Attention UNet）进一步提升了分割精度和泛化能力，适用于更复杂的应用场景。