U-Net是一种为图像分割而开发的卷积神经网络架构,尤其在生物医学领域应用广泛。它由奥拉夫·龙内贝格尔、菲利普·菲舍尔和托马斯·布罗克斯于2015年在弗莱堡大学提出,建立在早期全卷积网络的基础上。该架构以其对称的编码器-解码器设计和跳跃连接而著称,使其能够从相对较少的训练图像中生成精确的分割图。使用U-Net架构,在2015年的现代GPU上分割一张512 × 512的图像耗时不到一秒。
U-Net架构源于所谓的“全卷积网络”(FCN),由埃文·谢尔哈默、乔纳森·朗和特雷弗·达雷尔于2014年开发。其主要思想是通过连续的层来补充常规的收缩网络,其中池化操作被上采样算子取代。因此,这些层增加了输出的分辨率。随后的卷积层可以基于这些信息学习组装精确的输出。
网络架构
该网络由收缩路径和扩展路径组成,形成U形架构。收缩路径是一个典型的卷积网络,由重复的卷积操作组成,每次卷积后都跟随着一个修正线性单元(ReLU)和最大池化操作。在收缩过程中,空间信息减少,而特征信息增加。扩展路径通过一系列上卷积和与收缩路径中高分辨率特征的拼接,将特征和空间信息结合起来。
U-Net的一个重要修改是,在上采样部分有大量的特征通道,这使得网络能够将上下文信息传播到更高分辨率的层。因此,扩展路径与收缩部分大致对称,形成U形架构。网络仅使用每个卷积的有效部分,不包含任何全连接层。因此,网络可以以混合输入分辨率进行训练和运行。
为了预测图像边界区域的像素,通过镜像输入图像在图像边界处外推缺失的上下文。原始U-Net论文进一步提出了一种平铺策略,将大图像切割成重叠的瓦片并独立处理。这使得处理高分辨率图像成为可能,否则这些图像会超出可用的GPU内存。最近,基于感受野的U-Net模型在医学图像分割中也引起了关注。
训练与损失函数
U-Net使用随机梯度下降(SGD)进行端到端训练,采用逐像素损失函数,通常是交叉熵或其加权变体。能量函数通过对最终特征图进行逐像素soft-max计算,并结合交叉熵损失。作者引入了一个权重图,赋予接触对象边界附近的像素更高的重要性,这有助于网络学习在生物医学图像中分离单个实例。该权重图为每个训练图像预先计算,用于更重地惩罚边界区域的错误。
网络使用数据增强来增加有效训练集的大小,包括弹性变形、旋转和强度变化。这在标注数据稀缺的生物医学成像中尤为重要。原始实现使用高动量(0.99)和低学习率(0.00001)以确保稳定收敛。
在生物医学成像中的应用
U-Net在生物医学图像分割中有许多应用,例如在计算机断层扫描(CT)和磁共振成像(MRI)扫描中分割不同器官和身体系统。具体案例包括脑图像分割(BRATS)和肝脏图像分割(siliver07),以及蛋白质结合位点预测。U-Net实现也已在物理科学中找到用途,例如在材料显微图像分析中。
U-Net的变体也已应用于医学图像重建。以下是U-Net的一些变体和应用:
- 使用U-Net进行逐像素回归及其在全色锐化中的应用。
- 3D U-Net:从稀疏标注学习密集体积分割。
- TernausNet:使用在ImageNet上预训练的VGG11编码器的U-Net进行图像分割。
- 图像到图像转换以估计荧光染色。
- 在蛋白质结构结合位点预测中。
在扩散模型中的使用
U-Net架构也被用于扩散模型中的迭代图像去噪。该技术支撑了许多现代图像生成模型,如DALL-E、Midjourney和Stable Diffusion。在这些模型中,U-Net被训练来预测扩散过程中每一步添加到图像中的噪声,使模型能够逐步将随机噪声输入细化为连贯图像。U-Net中的跳跃连接在此特别有益,因为它们在去噪过程中保留了高频细节。
对语言模型的探索
U-Net也在被探索用于语言模型。分词不是单独的步骤,这使得模型更容易理解拼写,并同时向量化/分词更高层次的概念。这是一个新兴的研究领域,利用U-Net捕获多尺度特征的能力,可能为自然语言处理中占主导地位的Transformer架构提供替代方案。
历史与遗产
U-Net由奥拉夫·龙内贝格尔、菲利普·菲舍尔和托马斯·布罗克斯于2015年创建,并在论文“U-Net:用于生物医学图像分割的卷积网络”中报告。它是对FCN的改进和发展:埃文·谢尔哈默、乔纳森·朗和特雷弗·达雷尔(2014年)。“用于语义分割的全卷积网络”。该论文已成为深度学习领域被引用最多的论文之一,该架构已被改编用于无数超越生物医学成像的任务。
U-Net的成功启发了许多变体,包括具有不同编码器(如VGG或ResNet)、注意力机制和多尺度处理的变体。其影响扩展到机器学习的其他领域,包括生成式人工智能和计算机视觉。该架构也是医学影像挑战中的常见基线,并在TensorFlow和PyTorch等框架中被广泛实现。
实现
有多个U-Net的开源实现可用。J Akeret(2017年)的Tensorflow Unet是一个流行的实现。原始U-Net源代码可从德国弗莱堡大学计算机科学系的模式识别与图像处理小组获得。这些实现促进了U-Net在研究和工业中的广泛采用。