U-Net是一种用于图像分割的卷积神经网络架构,由Olaf Ronneberger、Philipp Fischer和Thomas Brox于2015年提出。该架构具有对称的编码器-解码器结构,并通过跳跃连接实现精确的定位,同时保留上下文信息。U-Net最初是为生物医学图像分割而开发的,但此后已被广泛应用于计算机视觉和生成建模等多个领域。其设计影响了众多后续架构,并至今仍是分割任务中的基准模型。
架构
U-Net架构由两条主要路径组成:收缩路径(编码器)和扩展路径(解码器),二者共同构成一个U形结构。收缩路径遵循典型的卷积网络布局,重复应用两个3x3卷积,每个卷积后接一个修正线性单元(ReLU)和一个步长为2的2x2最大池化操作以进行下采样。在每个下采样步骤中,特征通道的数量翻倍,使网络能够捕获越来越抽象的特征,同时降低空间分辨率。
相比之下,扩展路径对特征图进行上采样,随后进行一个2x2卷积,将特征通道数量减半。然后,该结果与收缩路径中相应裁剪过的特征图进行拼接,形成跳跃连接。这种拼接将编码器的高分辨率特征与解码器的上采样特征相结合,从而实现精确定位。之后,再应用两个3x3卷积,每个卷积后接一个ReLU。在最后一层,使用一个1x1卷积将特征向量映射到所需的类别数量。
U-Net的一个关键改进是上采样部分拥有大量的特征通道,这使得网络能够将上下文信息传播到更高分辨率的层。扩展路径与收缩路径大致对称,从而形成了标志性的U形结构。网络仅使用每个卷积的有效部分,即不进行填充,这导致每一步的空间维度略有减小。为了处理边界区域,缺失的上下文通过镜像输入图像来外推。原始论文还提出了一种针对大图像的平铺策略,即独立处理重叠的图块,然后将它们拼接在一起,从而能够处理超出GPU内存的高分辨率图像。
训练与数据增强
U-Net使用随机梯度下降进行训练,其损失函数将交叉熵与一种加权方案相结合,以强调边界像素。权重图是预先计算好的,用于分离同一类别的相邻物体,这在细胞经常重叠的生物医学分割中尤为有用。数据增强被用来有效增加训练集的大小,包括弹性变形和旋转,这些对于生物医学图像尤其重要,因为此类变化很常见。这种方法使得U-Net即使在训练样本相对较少的情况下也能实现高精度,这一点在原始论文中通过电子显微镜图像中神经元结构的分割得到了证明。
应用
U-Net已广泛用于生物医学图像分割,包括计算机断层扫描(CT)和磁共振成像(MRI)扫描中的器官和身体系统分割。具体实例包括BRATS挑战中的脑部图像分割和肝脏分割。在生物医学成像之外,U-Net还被用于诸如全色锐化中的逐像素回归等任务,以增强多光谱卫星图像的空间分辨率。它也被用于物理科学中材料显微图像的分析。U-Net的变体已被开发用于医学图像重建,并且该架构已被改编用于3D数据,例如用于从稀疏标注进行体积分割的3D U-Net。另一个值得注意的变体是TernausNet,它使用在ImageNet上预训练的VGG11编码器来提高分割任务的性能。
对生成模型的影响
U-Net架构在生成模型中也找到了重要应用,特别是在扩散模型中。在这些模型中,U-Net充当迭代图像去噪的主干网络,学习预测扩散过程中每一步添加到图像上的噪声。这项技术支撑了许多现代图像生成系统,包括DALL-E、Midjourney和Stable Diffusion。U-Net在保留细节的同时捕获全局上下文的能力使其非常适合这项任务。在此背景下,U-Net经常被修改以加入时间嵌入和注意力机制,从而处理变化的噪声水平并建模长距离依赖关系。
近期发展与扩展
U-Net启发了众多架构上的扩展和改进。注意力U-Net引入了注意力门控以聚焦于相关特征,而U-Net++(嵌套U-Net)则增加了密集跳跃连接以改善梯度流动。残差U-Net融入了残差块,以便于训练更深的网络。该架构也已与Transformer结合,例如TransUNet,它集成了一个Transformer编码器来捕获全局上下文。在扩散模型的背景下,带有注意力层的U-Net变体已成为标准。此外,人们还对基于感受野的U-Net模型用于医学图像分割产生了兴趣,旨在优化局部与全局信息之间的权衡。该架构至今仍是活跃的研究主题,人们不断努力提升其效率、准确性以及对新领域的适应性。
另见
参考文献
- Ronneberger, O., Fischer, P., & Brox, T. (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. In Medical Image Computing and Computer-Assisted Intervention (MICCAI).
- Shelhamer, E., Long, J., & Darrell, T. (2014). Fully Convolutional Networks for Semantic Segmentation. In CVPR.
- Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., & Ronneberger, O. (2016). 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation. In MICCAI.
- Iglovikov, V., & Shvets, A. (2018). TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation. arXiv preprint.
- Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. In NeurIPS.