U-Net学习笔记
简介
论文信息
U-Net出自2015年Olaf Ronneberger等人的论文 《U-Net: Convolutional Networks for Biomedical Image Segmentation
》,该论文可在IEEE explore中检索到。
U-Net提出背景
在过去几年里,深度卷积网络不仅广泛地应用于各种视觉识别任务,同时也达到了非常好的效果,尤其是在医学图像分割任务方面。然而由于医学图像可用训练集大小的限制以及考虑到网络大小,卷积网络难以取得更好的成果,幸运的是Hence, Ciresan等人提出了一种基于滑动窗口的网络,打破了这种局限性,该网络通过对某像素提供一个围绕其周围的局部区域(patch)并将该区域作为输入来预测该像素的所属类别。
滑动窗口网络
优点:
- 该网络可以进行局部定位(localize);
- 就“patch”而言,以它作为训练数据的数量要远多于训练图像的数量。
这刚好解决了卷积网络的瓶颈问题,但该网络同时也存在着缺点:
- 首先,它非常慢,因为必须为每个patch分别运行网络,并且由于patches之间的重叠导致了大量冗余;
- 其次,在定位精度和上下文语义信息使用之间存在折衷。 较大尺寸的patch需要更多的最大池化层,这会降低定位精度。而较小尺寸的patch又会使网络只能看到很少的上下文语义。
为了解决这一问题,人们开始考虑结合来自多层的特征,从而同时实现好的定位精度与上下文的使用,这就是要说的第二个网络,也是U-Net的基础网络——FCN。
FCN网络
FCN网络的主要思想是利用连续的层来辅助一个通常的收缩网络(contracting network,或者称为下采样网络),在这些层中,池化运算均被替换为上采样运算,因此这些层可以增加输出的分辨率。同时,为了得到更好的定位精度,FCN考虑融合来自收缩路径的高分辨率特征与上采样的输出,从而获取更精确的预测结果。(下图为FCN网络架构)基于FCN,Olaf Ronneberger等人修正并扩展该网络使得其在训练集很小的情况下也可以达到很好的分割效果,进而提出了U-Net网络。
网络架构
首先给出U-Net的网络结构图:
符号说明
蓝色箭头代表附有ReLU**函数的
3
∗
3
3*3
3∗3卷积操作;
红色箭头代表窗口为
2
∗
2
2*2
2∗2的最大池化操作;
绿色箭头代表
2
∗
2
2*2
2∗2反卷积操作;
浅绿色箭头表示最后一个
1
∗
1
1*1
1∗1卷积;
灰色箭头代表裁剪与融合(长连接);
蓝色框代表多通道的特征图。
网络组成
该网络分为左右两部分:
- 左半区域为contracting path,即下采样操作,其中重复赢哦用了两个 3 ∗ 3 3*3 3∗3卷积(未填充卷积),其后添加一个ReLU**函数与用于下采样的 2 ∗ 2 2*2 2∗2最大池化层,在每次下采样时都加倍特征通道数,达到的最低分辨率为 32 ∗ 32 32*32 32∗32;
- 右半区域为expansive path,即上采样操作。在上采样部分仍有大量的特征通道,这些特征通道允许网络将上下文信息传播到高分辨率层, 因此扩展路径或多或少地相对于收缩路径对称,进而产生u形结构。该区域每层都包含一个 2 ∗ 2 2*2 2∗2反的卷积操作用于上采样,该操作会减半特征通道。同时包含了与同维度层剪切后的特征图的融合操作,其后添加两个带有ReLU**函数的 3 ∗ 3 3*3 3∗3卷积。(由于我们使用了未填充卷积,因此每次卷积中都会丢失边界像素,因此有必要进行裁剪。)
- 在最后一层,使用
1
∗
1
1*1
1∗1卷积将每个64-分量特征向量映射到所需的类数。
该网络总共有23个卷积层。
该网络不包含任何的全连接层并且仅使用每个卷积的有效部分,即分割图仅包含像素信息,这使得上下文信息均可从输入图像中得到。进而,对任意大的输入图像,均可通过overlap-tile方法实现无缝分割,如下图所示。对于边界区域的预测,丢失的像素信息可以通过对输入图像的镜像来补充。这一方法对于使网络适应
网络训练
原文关于网络的训练主要强调了两点,一是由于训练集大小的限制需要使用数据增强技巧,二是生物医学图像分割中常见对于相接触的同一类的多个对象进行分割,因此文中使用了加权损失函数,使得接触细胞(touching cells)之间的分离背景标签在损失函数中获得较大权重。
该网络是在Caffe中利用动量为0.99的随机梯度下降算法实现的(SGD),其中每个batch一张图片,采用SoftMax+交叉熵损失。
加权损失函数
soft-max
其中
a
k
(
x
)
a_k(x)
ak(x)表示像素
x
x
x在特征通道
k
k
k上的**值(也就是像素
x
x
x对于类别
k
k
k的得分),
K
K
K表示类别数。容易看出,
p
k
(
x
)
p_k(x)
pk(x)是一个估计最大值函数,也就是说,对于取到最大**值
a
k
(
x
)
a_k(x)
ak(x)的
k
k
k而言,
p
k
(
x
)
≈
1
{p_k(x)}\approx1
pk(x)≈1,而对于其他
k
k
k而言,
p
k
(
x
)
≈
0
{p_k(x)}\approx0
pk(x)≈0。
权重 w ( x ) w(x) w(x)
文中给出了像素
x
x
x的权重map的计算方法:
其中
w
c
(
x
)
w_c(x)
wc(x)是用于平衡类别频率的权重图,
d
1
(
x
)
d_1(x)
d1(x)表示与
x
x
x最近的细胞之间的距离,
d
2
(
x
)
d_2(x)
d2(x)代表与
x
x
x第二近的细胞之间的距离。原文中设置
w
0
=
10
w_0=10
w0=10,
σ
=
5
\sigma=5
σ=5。
容易验证,当
x
x
x与附近细胞的距离越远,
w
0
e
x
p
(
)
w_0exp()
w0exp()越小,且逐渐趋于0,因此远离细胞的区域,其中像素的权值基本上都是一样的,也就是接近于
w
c
(
x
)
w_c(x)
wc(x)。相反,边界上的像素点的权重会随距离的增大而增大,进而容易区分同一类中相邻细胞的边界。
为了补偿某一类别的不同像素点的频次差异,用ground truth样例对权重图
w
(
x
)
w(x)
w(x)进行预先计算,可令网络学习区分同一类别不同实例之间更小的边界。如下图c和d:
权重
w
(
x
)
w(x)
w(x)通过高斯分布来初始化。
惩罚函数
惩罚函数的计算由
p
k
(
x
)
p_k(x)
pk(x)和
w
(
x
)
w(x)
w(x)决定:
其中
l
(
x
)
l(x)
l(x)为每个像素的真实类别。
数据增强
数据增强对于训练网络关于数据的不变性与鲁棒性非常重要,尤其是在医学图像数据集较少的情况下。文中就显微图像,需要使网络具有平移和旋转不变性、对于变形和灰度值变化的鲁棒性,尤其是在标注数据很少时,随机弹性形变(random elastic deformations)是训练分割网络的关键。
原文的处理方法是使用随机位移矢量在粗糙的
3
×
3
3\times3
3×3网格上生成平滑变形,位移矢量是从标准差为10的高斯分布中采样的,然后使用双三次方插值计算每个像素位移。
收缩路径最后的drop-out也起到了数据增强的作用。
实验
作者将U-Net应用到了三个分割任务。
电子显微镜记录中神经元分割
在该实验中,数据集由EM分割挑战提供,作者采用了三种指标对U-Net进行结果评价,实验结果如下:
光学显微镜中的细胞分割任务
该分割任务是ISBI细胞追踪挑战赛中的一部分,作者将U-Net应用到了2014和2015年的数据集,采用的评价指标为IOU (“intersection over union”) ,实验结果如下:
小结
无论是从网络架构还是实验结果来说,U-Net都是非常优秀的,它解决了医学图像分割领域中的一大难题,使得在可用数据集极大限制的情况下也能获得良好的分割效果,同时对于不同的分割任务都具有很好的效果。