UNet论文: 地址

UNet源代码: 地址

 

一、介绍

1.1 U-Net

很多分割网络都是基于FCNs做改进,包括这一章要讲的U-net。U-Net是一篇基本结构非常好的网络,由于网络结构像U型,所以叫Unet网络

U-Net不止应用于图像分割(尤其是医学图像分割),也用于自然图像生成的任务,所以它几乎是骨干(backbone)网络的标配

 

1.2 FCN和U-Net的区别

这里我引用了U-Net++作者的总结

在计算机视觉领域,全卷积网络(FCN)是比较有名的图像分割网络,医学图像处理方向,U-Net可以说是一个更加炙手可热的网络,基本上所有的分割问题,我们都会拿U-Net先看一下基本的结果,然后进行“魔改”。

U-Net和FCN非常的相似,U-Net比FCN稍晚提出来,但都发表在2015年,和FCN相比,U-Net的第一个特点是完全对称,也就是左边和右边是很类似的,而FCN的decoder相对简单,只用了一个deconvolution的操作,之后并没有跟上卷积结构第二个区别就是skip connection,FCN用的是加操作(summation),U-Net用的是叠操作(concatenation)。这些都是细节,重点是它们的结构用了一个比较经典的思路,也就是编码和解码(encoder-decoder)结构,早在2006年就被Hinton大神提出来发表在了nature上.

当时这个encoder-decoder结构提出的主要作用并不是分割,而是压缩图像和去噪声。输入是一幅图,经过下采样的编码,得到一串比原先图像更小的特征,相当于压缩,然后再经过一个解码,理想状况就是能还原到原来的图像。这样的话我们存一幅图的时候就只需要存一个特征和一个解码器即可。同理,这个思路也可以用在原图像去噪,做法就是在训练的阶段在原图人为地加上噪声,然后放到这个编码解码器中,目标是可以还原得到原图。

后来把这个思路被用在了图像分割的问题上,也就是现在我们看到的FCN或者U-Net结构,在它被提出的三年中,有很多很多的论文去讲如何改进U-Net或者FCN,不过这个分割网络的本质的拓扑结构是没有改动的。举例来说,去年ICCV上凯明大神提出的Mask RCNN. 相当于一个检测,分类,分割的集大成者,我们仔细去看它的分割部分,其实使用的也就是这个简单的FCN结构。说明了这种“U形”的编码解码结构确实非常的简洁,并且最关键的一点是好用。

此外, 由于UNet也和FCN一样, 是全卷积形式, 没有全连接层(即没有固定图的尺寸),所以容易适应很多输入尺寸大小,但并不是所有的尺寸都可以,需要根据网络结构决定, 在U-Net中, 池化是2x2的且是valid策略,即没有padding,因此要保证输入的图像在经过每一次池化的时候都要是边长偶数。所以要特别注意输入图像的尺寸。一个比较好的方法是从最小分辨率(分辨率就是feature map的尺寸
)出发沿收缩路径的反方向进行计算,得到输入图像的尺寸。

 

1.3 U-Net为什么在医学图像分割表现好?

U-Nnet模型在医学影像分割中表现好, 并且在以后许多优秀的医学图像分割网络中都是以它为basebone. 为了更好的分析原因, 不妨先看看医学图像有什么特点吧

1.3.1 医学影像的特点

尤其是相对于自然图像而言,医学图像具备以下几个特点:

  1. 图像语义较为简单、结构较为固定。我们做脑的,就用脑CT和脑MRI,做胸片的只用胸片CT,做眼底的只用眼底OCT,都是一个固定的器官的成像。由于器官本身结构固定和语义信息没有特别丰富,所以高级语义信息和低级特征都显得很重要(UNet的skip connection和U型结构就派上了用场)。
  2. 数据量少。医学影像的数据获取相对难一些,很多比赛只提供不到100例数据。所以我们设计的模型不宜过大,如果参数过多,很容易导致过拟合。
  3. 多模态。相比自然影像,医疗影像比较有趣和不同的一点是,医疗影像是具有多种模态的。以ISLES脑梗竞赛为例,其官方提供了CBF,MTT,CBV,TMAX,CTP等多种模态的数据。

    这就需要我们更好的设计网络去提取不同模态的特征feature。这里提供两篇论文供参考: Joint Sequence Learning and Cross-Modality Convolution for 3D Biomedical Segmentation(CVPR 2017) Dense Multi-path U-Net for Ischemic Stroke Lesion Segmentation in Multiple Image Modalities.

  4. 可解释性重要。由于医疗影像最终是辅助医生的临床诊断,所以网络告诉医生一个3D的CT有没有病是远远不够的,医生还要进一步的想知道,病灶在哪一层,在哪一层的哪个位置,分割出来了吗,能求体积嘛?同时对于网络给出的分类和分割等结果,医生还想知道为什么,也就是说从医生角度出发,更希望知道医疗AI的可解释性如何

    解释医学影像以Attention及CAM(class-activation-map)等可视化方法为主。

    https://www.zhihu.com/question/359161489?utm_source=qq

    https://www.zhihu.com/question/48224234

1.3.2 分析

  • 医学图像相比于普通图像,其复杂度是非常高的,灰度范围大,边界不清晰等特点, 并且人体内部结构相对固定,分割目标在人体图像中的分布很具有规律,语义简单明确且没有特别丰富, 而UNet结合了下采样时的低分辨率信息(提供物体类别识别依据)和上采样时的高分辨率信息(提供精准分割定位依据),此外还通过skip connection填补底层信息以提高分割精度.
  • 由于医学影像的数据获取相对难一些,数据量少, 所以如果模型的参数过多,很容易导致过拟合, 而U-Net模型大小较小,参数较少,因此比较适合.

PS: 分辨率就是feature map的尺寸

 

二、U-Net网络模型

Unet包括两部分:

  1. 特征提取部分,每经过一个池化层就一个尺度,包括原图尺度一共有5个尺度。
  2. 上采样部分,每上采样一次,就和特征提取部分对应的通道数相同尺度融合,但是融合之前要将其crop。这里的融合也是拼接。

该网络由收缩路径(contracting path)和扩张路径(expanding path)组成。其中,收缩路径用于获取上下文信息(context),即局部特征,扩张路径用于精确的定位(localization),且两条路径相互对称。

 

2.1 编码-解码结构(encoder-decoder)

U-Net和FCN它们的结构都用了一个比较经典的思路,也就是编码和解码(encoder-decoder)结构, 这种结构本来不是用在图像分割的,而是用在压缩图像和去噪声的[在本篇中的1.2有提到]

2.1.1 编码

前半部分是编码, 它的作用是特征提取(获取局部特征,并做图片级分类),这个降采样的理论意义是它可以增加对输入图像的一些小扰动的鲁棒性,比如图像平移,旋转等,减少过拟合的风险,降低运算量,和增加感受野的大小

从神经网络可视化那篇论文中《Visualizing and Understanding Convolutional Networks》能看出,底层的特征(如Layer1,2,3的输出)更偏向于组成图像的基本单元,如点,线,边缘轮廓,如下图:

而在高层抽象的特征(如Layer4,5)就更抽象,更近似于表示的是图像的语义信息,更像是一个区域了。如下图所示:

 

2.1.2 解码

后半部分是解码, 即利用前面编码的抽象特征来恢复到原图尺寸的过程, 最终得到分割结果

那么就有问题了,解码在将下采样(反卷积)数据恢复的时候,特征scale会发生变化,必然会有信息的丢失,这个时候, skip connention的作用就凸显出来了, skip connention起到了补充信息的作用.

 

2.2 skip connention 

skip connention 是为了实现特征融合的结构, 在FCN中用的是加操作(summation) ,而在U-Net中用的是叠操作(concatenation) ,这里也指的是复制和剪切(copy and crop).总之skip connention起到了补充信息的作用,让模型依赖更多的信息, 改善上采样时信息不足的问题, 从而提高分割精度.

总结下, 语义分割网络在特征融合时有2种办法:

  1. FCN式的逐点相加(即跳跃结构),对应caffe的EltwiseLayer层,对应tensorflow的tf.add()
  2. U-Net式的channel维度拼接融合(即copy and crop),对应caffe的ConcatLayer层,对应tensorflow的tf.concat()

在U-Net网络中可以发现灰色箭头的复制和剪切(copy and crop)操作,在同一层左边的最后一层要比右边的第一层要大一些,所以要想利用浅层的feature,就要进行一些剪切(以中心某个区域进行裁剪)。

 

此外, 可以看到, UNet融合了不同尺度的特征,同时跳级连接保证上采样恢复出来的特征不会很粗糙, 一个经验的解释就是跳级连接能够保证特征更加精细。这里推荐一篇论文,发表在AAAI 19 Workshop on Network Interpretability for Deep Learning的《Visualized Insights into the Optimization Landscape of Fully Convolutional Networks》,它通过可视化loss landscape对比发现这种跳级连接能够促使网络的minimizer更平坦,从而对新的数据敏感度更低,泛化能力更强

 

2.3 卷积

蓝色箭头代表3x3大小的卷积操作,并且stride是1,全程使用valid来进行卷积(卷积有三种模式,链接这里),因此,每个该操作以后,feature map的大小会减2。

valid卷积,即仅使用每个卷积的有效部分,不填充

此外,关于卷积还有两个重要的知识点:(具体可以参考我CNN的卷积知识)

  • 卷积层中的卷积核的通道只能设定为该层输入数据的通道,而大小可自定义(这里是3x3)
  • 卷积层中的卷积核的数量==输出特征图(output feature map)的数量(即输出数据的通道数)

 

2.4 池化

红色箭头代表2x2的max pooling操作,需要注意的是,此时的padding策略也是vaild

 

2.5 反卷积操作

绿色箭头代表2x2的反卷积操作,操作会将feature map的大小乘2。反卷积是上采样的一种方式.

 

2.6 最后一层1X1卷积

最后一层使用1 X 1大小的卷积核,将通道数降低至特定的数量(如像素点的类别数量). 这个过程是如何运算的呢?首先我们知道卷积层中的卷积核的通道只能设定为该层输入数据的通道,而大小可自定义, 并且卷积层中的卷积核的数量==输出特征图(output feature map)的数量(即输出数据的通道数) ,所以这里需要两个大小为1x1, 通道为64的卷积核进行卷积运算,才能得到两个通道的388*388大小的图.如下面的两个图所示(左边:多个卷积核的卷积运算模型图, 来自这里, 右变: 实际的卷积过程,步长S=1,填充=0)

最终的两个通道分别表示背景和目标(前景)两个类别

 

2.7 特征图上的像素级 softmax 值计算

接着模型的最后通过逐个像素地求其在这两张图像中该像素位置的最大数值描述(概率)作为该像素的分类, 因此产生了一张已经分割好的图片(output segmentation map),实现端到端.

那么这个概率怎么求, 论文中的损失函数首先是用了个pixel-wise softmax,就是每个像素对应的输出单独做softmax,也就是做了w*h个softmax, 

  • αk​(x) 代表在位置 x 处的像素在特征图中的第 k 层的激活值,也就是表示每一像素点(x)对应特征通道(k)的得分
  • K 是像素点的类别总数
  • pk(x)是类k的对像素点x的分类结果,就是概率

 

2.8 加权LOSS

  • 是每个像素的真实标签 
  • 是权重图,用来让某些像素更加重要 , 3.3会讲怎么获得权重图

 

三、创新(innovation)

下面的这些创新点并不适合所有任务,比如说对刚体进行分割,很难通过弹性变形对数据进行增强。

3.1 overlap-tile策略

在U-Net结构中,卷积全程都使用valid来进行卷积[2.3中提到],该卷积的特点就是没有padding, 即padding = 0,所以特征图(feature map)会越卷越小,它不会越卷越大,导致了最后输出的尺寸回不到原来输入的尺寸

原本以为U-Net是一个绝对对称的结构, 然而并不是,它左右两边feature map的尺寸不一致,原因就是因为valid卷积,只会越卷越小.如下图的层A,如果这个卷积有padding,那么它可以回到32x32,然后再上采样2x2后,层B就变为64x64,依次轮推最终实现结构的绝对对称

显然,如果生生把原始图像喂进来,最后输出的结果,会比原始图像小好多。我们希望输入跟输出尺寸一样,即不能强行scale up最后的输出,又不想给每层卷积加padding(连续对feature map加padding卷积,会使得padding进来的feature误差越来越大,因为越卷积,feature的抽象程度越高,就更容易受到padding的影响,所以为什么要用valid卷积,这是我参考这个作者总结的看法)

于是就诞生了重叠-切片(overlap-tile)策略, 该策略的思想是:对图像的某一块像素点(黄框内部分)进行预测时,需要该图像块周围的像素点(蓝色框内)提供上下文信息(context),以获得更准确的预测。简单地说, 就是在预处理中,对输入图像进行padding, 通过padding扩大输入图像的尺寸,使得最后输出的结果正好是原始图像的尺寸, 同时, 输入图像块(黄框)的边界也获得了上下文信息从而提高预测的精度

在U-Net原文中,用的镜像操作(mirror padding),实际上可以使用别的padding方法, 比如在这里做牙齿图像的segmentation,使用到了zero padding,即扩大的尺寸全用0(黑色)来填充。如下图所示:

因此, 使用Overlap-tile策略对数据进行预处理是有必要的,可对任意大的图像进行无缝分割. 而这种方法通常需要将图像进行分块(patch)的时候才使用

那么为什么要对图像分块(patch)再输入,而不是直接输入整张图像呢?因为内存限制,有的机器内存比较小,需要分块。然后再通过Overlap-tile 策略将分块图像进行预处理扩大输入尺寸.

总之,U-Net作者创新的Overlap-tile 策略,这种方法用于补全输入图像的上下信息,也可以解决由于现存不足造成的图像输入的问题

 

3.2 数据增强(data augmentation)

如果数据不够的话,尤其是医学影像的数据,可以用数据增强提高数据量.

U-Net作者的创新是使用了随机弹性变形进行数据增强,对数据进行增强,增加数据量,这允许网络可以学习到这种形变的不变形

此外,在收缩路径的末尾的drop-out层进一步暗示了数据增强,为什么要在收缩路径的末尾加drop-out层,而不是其它??

 

3.3 加权loss

作者通过预先计算权重图的来获得每一个像素在损失函数中的权值,这种方法补偿了训练数据每类像素的不同频率,并且使网络更注重学习相互接触的细胞间的边缘。分割边界使用形态学运算,特征图的计算方法如下:

  • ωc:Ω→R是用来平衡类频率的权重图 
  • d1:Ω→R表示到背景的某个像素点到最近细胞边界的距离 
  • d2:Ω→R表示到背景的某个像素点到第二近的细胞边界的距离

上图是用DIC(二次干涉对比)显微技术记录的玻璃上的 HeLa 细胞。其中图 (a) 是原始图像(gt);图 (b) 是基于 gt 的分割覆盖。其中不同的颜色表示不同的 HeLa 细胞示例。图 (c) 是生成的分割掩膜,其中白色部分是前景,黑色部分是背景;图 (d) 是像素级损失权重图,使得网络强制学习边缘像素。对于位与细胞接触部分的像素加大权重(红色的部分)

 

四、不足

U-Net++作者分析U-Net不足并如何做改进 :https://zhuanlan.zhihu.com/p/44958351

 

五、拓展

5.1 分块处理

医学图像是一般相当大,但是分割时候不可能将原图太小输入网络,所以必须切成一张一张的小patch,在切成小patch的时候,Unet由于网络结构的原因,因此适合用overlap的切图 [本篇的3.1 有讲],overlap部分可以为分割区域边缘部分提供文理等信息, 并且分割结果并没有受到切成小patch而造成分割情况不好。

5.2 UNet-family

https://github.com/ShawnBIT/UNet-family

 


参考文章:

https://zhuanlan.zhihu.com/p/31428783

https://blog.csdn.net/Formlsl/article/details/80373200

https://www.jianshu.com/p/14641b79a672

https://blog.csdn.net/hduxiejun/article/details/71107285

https://blog.csdn.net/qian99/article/details/85084686

https://www.zhihu.com/question/269914775?sort=created

https://blog.csdn.net/github_37973614/article/details/84559861

https://zhuanlan.zhihu.com/p/44958351

https://www.jianshu.com/p/b6a898753d29

https://www.zhihu.com/question/48224234

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐