摘要:人们普遍认为,深度网络的成功训练需要数千个带注释的训练样本。在本文中,我们提出了一种网络和训练策略,该策略依赖于大量使用数据增强来更有效地使用可用的注释样本。该体系结构包括用于捕获上下文的收缩路径和用于实现精确定位的对称扩展路径。我们表明,这样的网络可以从很少的图像进行端到端训练,并且在 ISBI 挑战中优于先前的最佳方法(滑动窗口卷积网络),用于分割电子显微堆栈中的神经元结构。使用在透射光显微镜图像(相差和 DIC)上训练的相同网络,我们在这些类别中以大幅优势赢得了 2015 年 ISBI 细胞跟踪挑战。而且,网络速度很快。在最新的 GPU 上分割 512x512 图像只需不到一秒。完整的实施(基于 Caffe)和经过训练的网络可在 http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net上获得。
引言
在过去的两年中,深度卷积网络在许多视觉识别任务中的表现优于现有技术,例如:虽然卷积网络已经存在了很长时间 ,但由于可用训练集的大小和所考虑网络的大小,它们的成功是有限的。 Krizhevsky 等人的突破是由于在具有 100 万训练图像的 ImageNet 数据集上对具有 8 层和数百万参数的大型网络进行了监督训练。从那时起,已经训练了更大更深的网络 。
卷积网络的典型用途是分类任务,其中图像的输出是单个类别标签。然而,在许多视觉任务中,尤其是在生物医学图像处理中,所需的输出应该包括定位,即应该为每个像素分配一个类标签。此外,数以千计的训练图像在生物医学任务中通常是遥不可及的。因此,Ciresan 等人在滑动窗口设置中训练了一个网络,通过提供该像素周围的局部区域(补丁)作为输入来预测每个像素的类标签。首先,这个网络可以本地化。其次,以块为单位的训练数据远大于训练图像的数量。由此产生的网络以大幅优势赢得了 ISBI 2012 的 EM 分割挑战。
显然,Ciresan 等人的策略有两个缺点。首先,它非常慢,因为网络必须为每个补丁单独运行,并且由于补丁重叠而存在大量冗余。其次,在定位精度和上下文的使用之间存在权衡。较大的补丁需要更多的 maxpooling 层,这会降低定位精度,而小补丁只允许网络看到很少的上下文。最近的方法提出了一个分类器输出,它考虑了多层的特征。可以同时进行良好的本地化和上下文的使用。
在本文中,我们建立在一个更优雅的架构之上,即所谓的“全卷积网络”。我们修改并扩展了该架构,使其适用于极少的训练图像并产生更精确的分割;参见图 1。中的主要思想是通过连续层补充通常的收缩网络,其中池化运算符被上采样运算符取代。因此,这些层增加了输出的分辨率。为了定位,来自收缩路径的高分辨率特征与上采样输出相结合。然后,连续的卷积层可以学习根据这些信息组装更精确的输出。
我们架构的一个重要修改是在上采样部分我们还有大量的特征通道,这允许网络将上下文信息传播到更高分辨率的层。因此,膨胀路径或多或少与收缩路径对称,并产生一个 U 形结构。该网络没有任何全连接层,仅使用每个卷积的有效部分,即分割图仅包含像素,输入图像中的完整上下文可用。该策略允许通过重叠平铺策略无缝分割任意大的图像(参见图 2)。为了预测图像边界区域中的像素,通过镜像输入图像来推断缺失的上下文。这种平铺策略对于将网络应用于大图像很重要,否则分辨率将受到 GPU 内存的限制。
至于我们的任务,可用的训练数据非常少,我们通过对可用训练图像应用弹性变形来使用过多的数据增强。这允许网络学习这种变形的不变性,而不需要在带注释的图像语料库中看到这些变换。这在生物医学分割中尤为重要,因为变形曾经是组织中最常见的变化,并且可以有效地模拟真实的变形。 Dosovitskiy 等人在无监督特征学习的范围内展示了数据增强对学习不变性的价值。
许多细胞分割任务中的另一个挑战是同一类别的接触对象的分离(参见图 3)。 为此,我们建议使用加权损失,其中接触单元格之间的分离背景标签在损失函数中获得较大的权重。
由此产生的网络适用于各种生物医学分割问题。在本文中,我们展示了 EM 堆栈中神经元结构分割的结果(一项持续的竞赛始于 ISBI 2012),我们的表现优于 Ciresan 等人的网络。此外,我们还展示了 2015 年 ISBI 细胞追踪挑战赛的光学显微镜图像中的细胞分割结果。在这里,我们在两个最具挑战性的 2D 透射光数据集上以大幅优势获胜。
网络结构
网络架构如图 1 所示。它由收缩路径(左侧)和扩展路径(右侧)组成。收缩路径遵循卷积网络的典型架构。它由重复应用两个 3x3 卷积(未填充卷积)组成,每个卷积后跟一个整流线性单元 (ReLU) 和一个步幅为 2 的 2x2 最大池化操作以进行下采样。在每个下采样步骤中,我们将特征通道的数量加倍。扩展路径中的每一步都包括特征图的上采样,然后是将特征通道数量减半的 2x2 卷积(“上卷积”),与来自收缩路径的相应裁剪特征图的连接,以及两个 3x3卷积,每个卷积后跟一个 ReLU。由于在每个卷积中丢失了边界像素,所以裁剪是必要的。在最后一层,使用 1x1 卷积将每个 64 分量特征向量映射到所需数量的类。网络总共有 23 个卷积层。为了实现输出分割图的无缝拼接(参见图 2),选择输入分块大小非常重要,这样所有 2x2 最大池化操作都会应用于具有偶数 x 和 y 大小的图层。
输入图像及其相应的分割图用于通过 Caffe 的随机梯度下降实现来训练网络。由于未填充卷积,输出图像比输入图像小一个恒定的边界宽度。为了最大限度地减少开销并最大限度地利用 GPU 内存,我们倾向于使用大输入图块而不是大批量大小,从而将批量减少为单个图像。因此,我们使用高动量 (0.99),以便大量先前看到的训练样本确定当前优化步骤中的更新。
能量函数由最终特征图上的逐像素 soft-max 结合交叉熵损失函数计算得出。 soft-max 定义为
其中 表示特征通道 k 在像素位置 x ∈ Ω 和 Ω ⊂ 处的激活。 K 是类别数, 是近似的最大函数。对于具有最大激活 的 k, ≈ 1,对于所有其他 k, ≈ 0。然后,交叉熵在每个位置惩罚 到 1 的偏差,使用
是每个像素的真实标签, 是我们引入的权重图,用于在训练中赋予某些像素更多的重要性。
我们预先计算每个地面实况分割的权重图,以补偿训练中来自特定类别的像素的不同频率数据集,并强制网络学习我们在接触单元之间引入的小分离边界(参见图 3c 和 d)。
使用形态学操作计算分离边界。然后将权重图计算为
其中 是平衡类别频率的权重图, 表示到最近单元格边界的距离, 表示到第二个最近单元格边界的距离。在我们的实验中,我们设置 和 像素
在具有许多卷积层和通过网络的不同路径的深度网络中,良好的权重初始化非常重要。否则,网络的某些部分可能会提供过多的激活,而其他部分则永远不会做出贡献。理想情况下,应调整初始权重,使网络中的每个特征图都具有近似单位方差。对于具有我们架构(交替卷积和 ReLU 层)的网络,这可以通过从标准偏差为 的高斯分布中提取初始权重来实现,其中 N 表示一个神经元的传入节点数。例如。对于 3x3 卷积和前一层的 64 个特征通道,N = 9 · 64 = 576
数据扩充
当只有很少的训练样本可用时,数据扩充对于教会网络所需的不变性和稳健性至关重要。在显微图像的情况下,我们主要需要移位和旋转不变性以及对变形和灰度值变化的鲁棒性。特别是训练样本的随机弹性变形似乎是训练具有很少注释图像的分割网络的关键概念。我们在粗糙的 3 x 3 网格上使用随机位移矢量生成平滑变形。
位移是从具有 10 个像素标准偏差的高斯分布中采样的。然后使用双三次插值计算每像素位移。收缩路径末尾的丢弃层执行进一步的隐式数据增强。
实验
我们演示了 u-net 在三种不同分割任务中的应用。第一项任务是在电子显微镜记录中分割神经元结构。图 2 显示了数据集和我们获得的分割的示例。我们提供完整的结果作为补充材料。该数据集由 EM 分割挑战赛提供,该挑战赛始于 ISBI 2012,并且仍在接受新的贡献。训练数据是一组 30 张图像(512x512 像素),来自果蝇一龄幼虫腹神经索 (VNC) 的连续切片透射电子显微镜。每张图像都带有相应的完整注释的细胞(白色)和细胞膜(黑色)地面实况分割图。测试集是公开可用的,但其分割图是保密的。可以通过将预测的膜概率图发送给组织者来获得评估。评估是通过在 10 个不同级别对地图进行阈值处理并计算“扭曲误差”、“兰德误差”和“像素误差”来完成的。
u-net(对输入数据的 7 个旋转版本进行平均)在没有任何进一步预处理或后处理的情况下实现了 0.0003529 的扭曲误差(新的最佳分数,请参见表 1)和 0.0382 的随机误差。
这明显优于 Ciresan 等人的滑动窗口卷积网络结果,后者的最佳提交具有 0.000420 的扭曲误差和 0.0504 的随机误差。就随机误差而言,该数据集上唯一表现更好的算法使用高度数据集特定的后处理方法 1 应用于 Ciresan 等人的概率图。
我们还将 u-net 应用于光学显微图像中的细胞分割任务。此分割任务是 2014 年和 2015 年 ISBI 细胞跟踪挑战 [10,13] 的一部分。第一个数据集“PhC-U373”2 包含胶质母细胞瘤-星形细胞瘤 U373 细胞,位于通过相差显微镜记录的聚丙烯酰亚胺基质上(参见图 4a、b 和补充材料)。它包含 35 个部分注释的训练图像。在这里,我们实现了 92% 的平均 IOU(“并集交集”),这明显优于 83% 的第二佳算法(见表 2)。第二个数据集“DIC-HeLa”3 是通过微分干涉对比 (DIC) 显微镜记录的平板玻璃上的 HeLa 细胞(参见图 3、图 4c、d 和补充材料)。它包含 20 个部分注释的训练图像。在这里,我们实现了 77.5% 的平均 IOU,这明显优于第二好的算法 46%。
结论
u-net 架构在非常不同的生物医学分割应用程序上实现了非常好的性能。得益于弹性变形的数据增强,它只需要很少的注释图像,并且在 NVidia Titan GPU (6 GB) 上的训练时间非常合理,仅为 10 小时。我们提供完全基于 Caffe的实现和训练有素的网络 4。我们确信 u-net 架构可以轻松应用于更多任务。