论文阅读——Recurrent U-Net for Resource-Constrained Segmentation

论文阅读之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框架。