相关信息
论文地址:
Conference paper (发表于DLMIA):UNet++: A Nested U-Net Architecture for Medical Image Segmentation
Journal paper (发表于IEEE TMI):UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation
代码地址:
Code:github.com/UNetPlusPlus
注
:本文是一篇阅读笔记。作者知乎解读:研习U-Net . 建议先看这篇知乎博文!
我一开始的打算是在我这篇博文上做简单的重点笔记,但是作者写的太好了,每个环节环环相扣,所以后来打算大部分地转载作者这篇文章,便于自己查看。
1. 铺垫
在计算机视觉领域,全卷积网络(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上.
当时这个结构提出的主要作用并不是分割,而是压缩图像和去噪声。输入是一幅图,经过下采样的编码,得到一串比原先图像更小的特征,相当于压缩,然后再经过一个解码,理想状况就是能还原到原来的图像。这样的话我们存一幅图的时候就只需要存一个特征和一个解码器即可。同理,这个思路也可以用在原图像去噪,做法就是在训练的阶段在原图人为的加上噪声,然后放到这个编码解码器中,目标是可以还原得到原图。
后来把这个思路被用在了图像分割的问题上,也就是现在我们看到的U-Net结构,在它被提出的三年中,有很多很多的论文去讲如何改进U-Net或者FCN,不过这个分割网络的本质的拓扑结构是没有改动的。举例来说,去年ICCV上凯明大神提出的Mask RCNN 相当于一个检测,分类,分割的集大成者,我们仔细去看它的分割部分,其实使用的也就是这个简单的FCN结构。说明了这种“U形”的编码解码结构确实非常的简洁,并且最关键的一点是好用。
简单的过一下这个网红结构,我们先提取出它的拓扑结构,这样会比较容易分析它的实质,排除很多细节的干扰。
输入是一幅图,输出是目标的分割结果。继续简化就是,一幅图,编码(或者说降采样),然后解码(也就是升采样),然后输出一个分割结果,如下图所示。根据结果和真实分割的差异,反向传播来训练这个分割网络。
我们可以说,U-Net里面最精彩的部分就是这三部分:
- 下采样
- 上采样
- skip connection
大部分分割网络的拓扑结构:
2. 提出问题
大家都在这个U形拓扑结构做什么样的改进呢?如果让你在这个经典的结构基础上改进,你会去关注哪些点呢?
首先一个问题是:
问题1:四次降采样是不是必须?要多深合适?
这里作者想强调的一点是:很多论文给出了他们建议的网络结构,其中包括非常多的细节,比如用什么卷积,用几层,怎么降采样,学习率多少,优化器用什么,这些都是比较直观的参数,其实这些在论文中给出参数并不见得是最好的,所以关注这些的意义不大。
一个网络结构,我们真正值得关注的是它的设计传达了什么信息。就拿U-Net来说,原论文给出的结构是原图经过四次降采样,四次上采样,得到分割结果,实际呢,**为什么四次?**就是作者喜欢呗,或者说当时作者使用的数据集,四次降采样的效果好,我们也可以更再专业一点,四次降采样的接受域或者感受野大小正合适处理图像。或者四次降采样比较适合输入图像的尺寸等等,理由一堆,但是你们真的相信吗?不见得吧。
答1:四次降采样不是必须,要根据任务的复杂程度来设计网络深度。
2017年在CVPR上发表的一个名叫PSPNet的分割网络,你会发现,好像整体的架构和U-Net还是像的,只是降采样的数目减小了,当然,他们也针对性的增强了中间的特征抓取环节的复杂性。
要是你还觉得这个工作还不够说明4次降采样不是必须的话,还可以看看Yoshua Bengio关于图像分割的论文,他们提出的结构,名叫提拉米苏。
也是U形结构,但他们就只用了三次降采样。所以到底要多深,是不是越深越好?这是一个开放的问题。
关于到底要多深这个问题,还有一个引申的问题:
问题2:降采样对于分割网络到底是不是必须的?
问这个问题的原因就是,既然输入和输出都是相同大小的图,为什么要折腾去降采样一下再升采样呢?
比较直接的回答当然是降采样的理论意义:
答案2:从理论意义上:它可以增加对输入图像的一些小扰动的鲁棒性,比如图像平移,旋转等,减少过拟合的风险,降低运算量,和增加感受野的大小。升采样的最大的作用其实就是把抽象的特征再还原解码到原图的尺寸,最终得到分割结果。
这些理论的解释都是有道理的。
在作者的理解中,对于特征提取阶段,浅层结构可以抓取图像的一些简单的特征,比如边界,颜色,而深层结构因为感受野大了,而且经过的卷积操作多了,能抓取到图像的一些说不清道不明的抽象特征,讲的越来越玄学了,总之,浅有浅的侧重,深有深的优势。那我就要问一个比较犀利的问题:
问题3:既然浅层特征和深层特征都很重要,U-Net为什么只在4层以后才返回去,即降采样16倍,也就是只去抓深层特征。
答3: 这个问题的解答其实和第一问很相似。并不是非要降低16倍,而是根据自身问题的难度选择降低的倍数。
一个可能有用的建议:数据集越复杂分割难度越大,试试用更深层的特征。
作者为了说明白这个问题本身,结合图片给出了说明:如下图所示,既然 X 1 , 0 X^{1,0} X1,0、 X 2 , 0 X^{2,0} X2,0、 X 3 , 0 X^{3,0} X3,0、 X 4 , 0 X^{4,0} X4,0所抓取的特征都很重要,为什么非要降到 X 4 , 0 X^{4,0} X4,0层才开始上采样回去呢?
然后作者设计了对比实验,对深度不一的U-Net结构进行实验,对比其分割表现:
先不要看后两个UNet++,就看这个不同深度的U-Net的表现,我们可以看出,不是越深越好吧,它背后的传达的信息就是,不同层次特征的重要性对于不同的数据集是不一样的,并不是说我设计一个4层的U-Net,就像原论文给出的那个结构,就一定对所有数据集的分割问题都最优。
从上图可以看出U-Net++分割效果比U-Net效果好,但是这个网络增加了多少参数?
这个问题非常的尖锐,实际上是需要设计实验来回答的,如何设计这个实验呢?作者的做法是强行增加U-Net里面的参数量,让它变宽,也就是增加它每个层的卷积核个数。由此,设计了一个叫wide U-Net的参考结构,先来看看UNet++的参数数量是9.04M,而U-Net是7.76M,多了差不多16%的参数,所以wide U-Net在作者设计时就让它的参数比UNet++差不多,并且还稍微多一点点,来证明并不是无脑增加参数量,模型效果就会好。
作者本人是认为这个回击力度并不大,因为这样增加参数说心里话有点敷衍,应该能找到更好的对比方法。尽管有不足,我们先来看看结果:
根据现有的结果,作者的总结是,单纯的把网络变宽,把参数提上去对效果的提升并不大,如何能把参数用在刀刃上是很重要的。那么UNet++这种设计就是将参数用在刀刃上的一个方案。
我们在回来看这个UNet++,对于这个主体结构,作者在论文中给出了一些点评,说白了就是把原来空心的U-Net填满了,优势是可以抓取不同层次的特征,将它们通过特征叠加的方式整合,不同层次的特征,或者说不同大小的感受野,对于大小不一的目标对象的敏感度是不同的,比如,感受野大的特征,可以很容易的识别出大物体的,但是在实际分割中,大物体边缘信息和小物体本身是很容易被深层网络一次次的降采样和一次次升采样给弄丢的,这个时候就可能需要感受野小的特征来帮助。另一个解读就是如果你横着看其中一层的特征叠加过程,就像一个去年很火的DenseNet的结构,非常的巧合,原先的U-Net,横着看就很像是Residual的结构,这个就很有意思了,UNet++对于U-Net分割效果提升可能和DenseNet对于ResNet分类效果的提升,原因如出一辙,因此,在解读中作者也参考了Dense Connection的一些优势,比方说特征的再利用等等。
3. 方法和实验验证
刚刚在讲这里的时候留了一个伏笔,说这个结构在反向传播的时候中间部分会收不到过来的梯度,如果只用最右边的一个loss来做的话。
一个非常直接的解决方案就是深监督,也就是deep supervision。这个概念不是新的,有很多对U-Net对改进论文中也有用到,具体的实现操作就是:
方法1:在图中 X 0 , 1 X^{0,1} X0,1、 X 0 , 2 X^{0,2} X0,2、 X 0 , 3 X^{0,3} X0,3、 X 0 , 4 X^{0,4} X0,4后面加一个1x1的卷积核,相当于去监督每个level,或者每个分支的U-Net的输出。
作者给大家提供三个用Deep Supervision的结构来比较,一个就是上面这个,第二个是:
方法2:加在UC Berkeley提出的结构上,结构如下图所示:
最后一个是:
方法3:加在UNet++上
请问是你认为哪个会更棒?这是一个开放问题,我这里不做展开讨论,作者论文中使用的是最后一个。当它配合上这样一个填满的U-Net结构时,带来其中一个非常棒的优势。
问题:想一想,如果在训练过程中在各个level的子网络中加了这种深监督,可以带来怎样的好处呢?
答:两个字:剪枝。
同时引出三个问题:
问题:
(1)为什么UNet++可以被剪枝?
(2)如何剪枝?
(3)好处在哪里?
答:
(1) 被剪枝其实就是因为网络中间层就能得到满意的结果。
(2)减去最外层的一排,或减去最外层若干排(根据实验结果来确定具体减多少)。
(3)减少参数的同时还能保证一定的性能。
来看看为什么可以剪枝,这张图特别的精彩。关注被剪掉的这部分,你会发现,在测试的阶段,由于输入的图像只会前向传播,扔掉这部分对前面的输出完全没有影响的,而在训练阶段,因为既有前向,又有反向传播,被剪掉的部分是会帮助其他部分做权重更新的。这两句话同样重要,再重复一遍,测试时,剪掉部分对剩余结构不做影响,训练时,剪掉部分对剩余部分有影响。这意味什么?
因为在深监督的过程中,每个子网络的输出都其实已经是图像的分割结果了,所以如果小的子网络的输出结果已经足够好了,我们可以随意的剪掉那些多余的部分了。
来看一下上面这个图,为了定义的方便起见,我们把每个剪完剩下的子网络根据它们的深度命名为UNet++ L1,L2,L3,L4,后面会简称为L1~L4。最理想的状态是什么?当然是L1喽,如果L1的输出结果足够好,剪完以后的分割网络会变得非常的小。
这里作者又提出了两个问题:
问题1:为什么要在测试的时候剪枝,而不是直接拿剪完的L1,L2,L3训练?
问题2:怎么去决定剪多少?
答1:对于为什么要在测试的时候剪枝,而不是直接拿剪完的L1,L2,L3训练,作者的解释其实上面写了,剪掉的部分在训练时的反向传播中是有贡献的,如果直接拿L1,L2,L3训练,就相当于只训练了不同深度的U-Net,最后的结果会很差。
答2:如何去决定剪多少,还是比较好回答的。因为在训练模型的时候会把数据分为训练集,验证集和测试集,训练集上是一定拟合的很好的,测试集是我们不能碰的,所以我们会根据子网络在验证集的结果来决定剪多少。所谓的验证集就是一开始从训练集中分出来的数据,用来监测训练过程用的。
讲完了思路,我们来看结果:
先看看L1~L4的网络参数量,差了好多,L1只有0.1M,而L4有9M,也就是理论上如果L1的结果我是满意的,那么模型可以被剪掉的参数达到98.8%。不过根据我们的四个数据集,L1的效果并不会那么好,因为太浅了嘛。但是其中有三个数据集显示L2的结果和L4已经非常接近了,也就是说对于这三个数据集,在测试阶段,我们不需要用9M的网络,用半M的网络足够了。
回想一下一开始我提出的问题,网络需要多深合适? 这幅图就一目了然了。
网络的深度和数据集的难度是有关系的,这四个数据集当中,第二个,也就是息肉分割是最难的,大家可以看到纵坐标,它代表分割的评价指标,越大越好,其他都能达到挺高的,但是唯独息肉分割只有在30左右,对于比较难的数据集,可以看到网络越深,它的分割结果是在不断上升的。对于大多数比较简单的分割问题,其实并不需要非常深,非常大的网络就可以达到很不错的精度了。
上图中,横坐标代表的是在测试阶段,单显卡12G的TITAN X (Pascal)下,分割一万张图需要的时间。我们可以看到不同的模型大小,测试的时间差好多。如果比较L2和L4的话,就差了三倍之多。
对于测试的速度,用下面这一幅图会更清晰。
我们统计了用不同的模型,1秒钟可以分割多少的图。如果用L2来代替L4的话,速度确实能提升三倍。
剪枝应用最多的就是在移动手机端了,根据模型的参数量,如果L2得到的效果和L4相近,模型的内存可以省18倍。还是非常可观的数目。
关于剪枝的这部分我认为是对原先的U-Net的一个很大的改观,原来的结构过于刻板,并且没有很好的利用不用层级的特征。
再次放上U-Net++网络结构:
简单的总结一下:
UNet++的第一个优势:精度的提升,这个应该它整合了不同层次的特征所带来的;
UNet++的第二个优势:灵活的网络结构配合深监督,让参数量巨大的深度网络在可接受的精度范围内大幅度的缩减参数量。
附录. 模型网络改进建议
注:这是给自己的一些启示和想法
(1)感受野大的特征,可以很容易的识别出大物体的,但是在实际分割中,大物体边缘信息和小物体本身是很容易被深层网络一次次的降采样和一次次升采样给弄丢的
(2)降采样的理论意义:它可以增加对输入图像的一些小扰动的鲁棒性,比如图像平移,旋转等,减少过拟合的风险,降低运算量,和增加感受野的大小。升采样的最大的作用其实就是把抽象的特征再还原解码到原图的尺寸,最终得到分割结果。
(3)网络的深度和数据集的难度是有关系的,这上面四个数据集当中,第二个,也就是息肉分割是最难的,看到纵坐标,它代表分割的评价指标,越大越好,其他都能达到挺高的,但是唯独息肉分割只有在30左右,对于比较难的数据集,可以看到网络越深,它的分割结果是在不断上升的。对于大多数比较简单的分割问题,其实并不需要非常深非常大的网络就可以达到很不错的精度了。