Convolutional Networks for Biomedical Image Segmentation
Olaf Ronneberger, Philipp Fischer, and Thomas Brox
University of Freiburg, Germany, 2015
用于语义分割的U-Net,后面的很多SOTA语义分割架构都采用了这个架构。
1. Introduction
深度学习一直受限于数据集大小和网络大小。
我们的一个重要修改是在上采样部分也有很多的特征通道,这让网络可以从上下文信息向更高分辨率层传播。因此扩张路径或多或少地与压缩路径是对称的,形成了一个U型。网络不含任何全连接层,这使得任意大小的图片都可以用重叠平铺的方式(图2)来进行无缝分割。为了预测图片边缘的像素,缺失的上下文通过对输入图片的镜像进行推断。这一策略对将网络应用于大图很重要,否则分辨率会被显存大小所限制。
我们的训练数据非常少,因此我们通过弹性形变进行了大量数据增广。因为组织的常见变化就是弹性形变。
细胞分割的另一挑战是将接触的同类物体分开,如图3。我们为此提出了加权loss,在接触的细胞间的背景标签会在loss函数中有很大权重。
得到的网络能应用于多个生物医学分割问题。
2. Network Architecture
网络结构见图1。它由收缩路径(左侧)和扩展路径(右侧)组成。收缩侧是典型的卷积网络架构。它由重复的两个$3\times3$卷积(无padding)组成,每个后跟一个ReLU和一个步长为2的$2\times 2$max pooling。在每次下采样时我们都会让通道数加倍。扩展侧中每一步由一个上采样和一个$2\times2$卷积(”上卷积”)组成,对通道数减半,从收缩侧对应的特征图拷贝、裁切并连接,以及两个各跟一个ReLU的$3\times3$卷积。裁切的原因是边缘像素在每次卷积都会损失。最终层会使用$1\times1$卷积将每个64位的向量转换为需要的类数量。整个网络共有23个卷积层。
为了任意地堆叠输出的分割图,需要小心选择输入尺寸,在所有进行$2\times2$ max pooling时长宽都是奇数。
3. Training
使用图片输入和分割图来对网络进行训练。由于unpadded卷积,输出的图片尺寸会比输入小常量的边宽。为了最小化间接花费overhead,最大化利用显存,我们倾向于使用更大的块大小而不是更大的batch size,因此将batch退化到单张图片。于是我们使用了大的动量(0.99),使上一张的样本会决定当前步的更新。
Energy函数是由一个在最终特征图的逐像素soft-max和交叉熵损失函数组成。soft-max定义如下:$p_k(\mathbf x) = \left. \exp (a_k(\mathbf x)) \middle / \right. \left (\sum _{k’=1}^K \exp (a_{k’}(\mathbf x)) \right )$,其中$a_k(\mathbf x)$记做在像素位置$\mathbf x$的第k个通道的激活值。K是类别数,$p_k(\mathbf x)$是近似的maximum-function。例如,当$a_k(\mathbf x)$是最大的激活值时,$p_k(\mathbf x) \approx 1$,其余的k的$p_k(\mathbf x) \approx 0$。交叉熵通过以下公式对每个位置的$p_{l(\mathbf x)} (\mathbf x)$与1的偏移进行惩罚:
其中$l : \Omega \to \{ 1, …,K\}$是每个像素的true标签,$w: \Omega \to \mathbb R$是权重map,我们使用它让某些像素在训练中更重要。
我们为每个gt分割预计算权重图,来补偿训练集中不同类别的评率不同,强迫网络学习相邻细胞的小分隔边。
分隔边使用形态学morphological操作计算。权重图计算如下:
其中$w_c: \Omega \to \mathbb R$是用于平衡类频率的权重图,$d_1: \Omega \to \mathbb R$记做与最近细胞的边界距离,$d_2: \Omega \to \mathbb R$是与到第二近的细胞边界距离。我们的实验中,$w_0 = 10$, $\sigma \approx 5$像素。
有许多卷积层和不同的路径的网络,一种好的权重初始化方式非常重要。否则,部分网络会给出过多的激活,而其他的部分从未贡献。理想情况应使网络中的每个特征图有着近似的单元方差。对于我们这样的架构(交替alternating的卷积和ReLU层),可以从有着标准差$\sqrt{2/N}$的高斯分布获得初始权重,N是一个神经元的输入节点数量。例如对于一个$3\times3$卷积层,前一层有64个通道,$N=9\cdot 64=576$。
3.1 Data Augmentation
对于显微镜图片,我们主要需要偏移和旋转不变性,变形的健壮性以及灰度值的变化。尤其是随机弹性变形似乎是在训练集有限时训练一个分割网络的关键点。我们在一个粗糙的$3\times3$grid上使用随机向量替换来生成光滑的变形。替换采样自标准差为10像素的高斯分布。接着使用双三次插值bicubic interpolation进行逐像素的替换。收缩路径上的Dropout层也隐式地进行了进一步的数据增广。
4. Experiments
我们将U-Net应用到了多个任务上,首先是电子显微镜记录的神经元结构分割。图2是一个样本以及我们的分割结果。数据集是EM segmentation challenge。训练集是30张512像素的图片。