论文阅读之Recurrent U-Net for Resource-Constrained Segmentation
解决在资源受限平台上使用Recurrent UNet做分割任务
from ICCV2019
# 摘要
目前许多性能较好的模型都是基于深层次的网络、大规模的数据集,在GPU上跑起来就很慢,本文则提出了一种新的Recurrent UNet即保留了UNet的紧凑型又在一些benchmark上获得了很好的性能提升,如手部分割、视网膜血管分割以及道路分割任务中。
Section I Introduction
我们发现一些性能较好的予以分割框架都是基于深层次网络、高分辨率特征、大规模数据集以及需要预训练,比如DeepLab系列、PSPNet和RefineNet,往往对GPU要求很高,推断事件也很长,但对于一些资源有限的场合或者需要实时分割的场合不太适用。所需要的大型数据集也不适合医学图像的分割,在医学图像分割领域主流还是UNet系列
因此本文基于紧凑的UNet提出了recurrent UNet用于资源有限的平台,而且比现在的一些benchmark性能上均有提升;参量方面仅有0.3million,比现有的基于ResNet-101的DeepLab和RefineNet都要精简(他们的参量分别是40million、118million)
Unet经过编解码网络后直接又输入得到分割图像,而本文的recurrent体现在分割任务和网络内部状态的迭代循环,对UNet中间层进行迭代优化,就像人基于注意力积累知识的过程一样,而且还可以灵活选择哪些状态进行recurrent。非常适合训练数据有限/资源有限的场景,推断起来也符合实时性要求(55 fps)而且可以为了提速减少迭代次数(虽然一定程度损失精度)
Recurrent UNet总结:
基于UNet
引入recurrent
适合资源有限、数据集规模不大、实时性好
# Section II Related Work
紧凑语义分割模型
主要有两类,一类是基于编解码结构(encoder-decoder)的网络如医学图像分割领域最常见的UNet;另一类基于多分支结构(multi-branch)将不同层次分辨率的特征进行融合,这对融合方式要求比较高,因为要权衡速度和精度;UNet通过skip connection跃层连接的方式比这简单多了,也是UNet受欢迎的原因。
循环分割网络:
这一思想在AutoContext中被首次提出(还有循环随机森林)。
整体循环:比如将UNet输出的mask结合原图再次作为输入体现循环的思想,每次迭代输入patch的分辨率逐渐精细,用来产生更加精细的分割图。
内部循环:
另一种迭代思想的体现是设置一个循环单元(recurrent unit),比如将UNet中的一些层通过一个GRU门控单元进行循环迭代。
Section III Method
Part A Recurrent UNet整体架构
我们的框架整体和UNet基础框架一致的,都是基于编解码网络——跃层连接,但由于是轻量级网络,所以通道数上和原始UNet不太一致,conv1通道数为8随后每次maxpool后通道数double一倍,decoder部分采用转置卷积,目的是为了比双线性插值获得更好的表达能力。
循环的思想体现在(1)分割结果的循环迭代(2)中间层的循环迭代
(1)就是将前一状态的输出(hidden tensor)和原始图像(input image)以级联的方式作为下一次循环的输入
(2)的实现则是将UNet中的部分层用一个循环单元(recurrent unit)替代,本文从layer=3开始作为一个循环Unit,如果l=0意味着整个UNet都处于一个循环单元中。
接下来讨论recurrent unit的具体结构-DRU&SRU两种,区别在于门控的数目不同。
UNet & R-UNet
Part B Recurrent Unit
Dual-gated Recurrent Unit(DRU)
灵感来源于GRU门控单元,将UNet某些编解码层替换为recurrent unit,将上一状态的激活值+前层的输出作为输入来进行更新,设置reset gate和update gate两个门控单元决定状态是否更新。
Single-gated Recurrent Unit(SRU)
DRU由于包含了多层编解码结构的循环,对memory需求比较大,因此为了减少内存开销设计了SRU单元,减少了一个门控(reset gate)
DRU & SRU
Part C Training
通过逐像素的交叉熵损失函数计算loss值,每一次迭代引入监督机制,通过alpha这个参数来决定每一次迭代的权重(比如本文迭代三次,对权重更新的比例是一样的,也可以给不同迭代次数设置不同的更新权重,比如0.3 0.3 0.4给最后一次迭代赋予更大的更新权重。)
Section IV Experiments
实验测试了recurrent两种版本的效果用于手势分割、视网膜血管分割以及道路分割。
Part A 数据集
Hands:即在常规的手势数据集EYTH、GTEA\EgoHand\HOF上进行了测试(4800张)也在自己准备的大型数据集KBH Dataset上进行了测试(打字的动作)。数据集划分:2:2:6=train/val/test,这样做是不用过多测试数据评估学习能力和泛化性。
Retina Vessels:视网膜血管分割用的就是DRIVE数据集了共40张图片(20train:20test),其中33张正常7张为糖尿病视网膜早期病变。
Roads:
Massachusetts Roads dataset1108train:49test;
Urban landscapes:Cityspaces 提供5000张高分辨率(1024*2016)张图片,划分情况:2975:500:1525=train:val:test
Part B 实验设定
将本文以UNet为骨架的DRU/SRU与RefineNet、ICNet、Deeplab进行对比;对比的还有:
UNet-B UNet-G分别是使用BN和group Normalization的;
Rec-Last:是在网络最后加一个recurrent unit
Rec-Middle:Unet bottleneck放Recurrent Unit
Rec-Simple:将分割产生的mask与原图再次输入网络用一种recursive refinement的方式
还有将UNet的encoder部分用预训练过的VGG-16替换。-heavy one
与此对应的light-one拥有更少的参数量以及从头开始训练
Hands:提供了数据上和可视化了分割结果,可以看到在light轻量级系列使用recurrent -unit是比one-shot的ICNet性能更好,而且使用DRU单元的性能往往比SRU性能更佳;在heavy系列的对比中可以看到R-Unet在5个数据集中4个都比RefineNet要好,出了一个小规模的HOF数据集(198trainging images),而且RefineNet是高度依赖于ResNet-101预训练的基础上,而R-UNet是从头开始训练的;
实验还测试了将UNet的encoder部分用预训练过的VGG-16代替,虽然在HOF等性能有所提升,但参数量可是原始R-UNet的100倍,帧率分别为18fps和61fps
总结就是参量少 实时性好
Retina Vessel:在DRIVE数据集做血管分割的结果显示R-UNet甚至比预训练过得DRU-VGG-16以及DeepLab系列都要优秀,主要原因可能是因为DRIVE数据集较小,深层次的网络更容易过拟合,以及DRIVE数据集和ImageNet的图像差别还是蛮大的,因此预训练不预训练的提升效果不大。
医学图像分割:R-UNet由于recurrent的存在 对小规模数据集也很适合,而且对预训练的backbone要求并不是必须的。
Roads:这一部分使用R-UNet比其他baseline model效果都要好,尤其是使用DRU获得了8%的提升;还发现如果将通道数double会进一步提升acc,原始的R-UNet channel=8,16,32。
Section V Conclusion
本文基于紧凑的UNet网络结合循环迭代的思想提出Recurrent-UNet(R-UNet)结构,核心思想就是将UNet中部分encoder-decoder层组合成Recurrent-unit循环迭代,主要特点是通道少,参量少,轻量级以及实时性好。
本文的其他贡献还有就是KBH数据集,比其他现有的手势数据集规模都要大。
未来优化方向:换到其他backbone上测试recurrent-unit的有效性;尝试不同层次的recurrent框架。