论文阅读——Non-local U-Nets for Biomedical Image Segmentation

论文阅读之非局部UNet模型用于医学图像分割
Non-local U-Nets for Biomedical Image Segmentation
fromAAAI 2020
Code

摘要：

UNet的编解码结构广泛用于医学图像分各领域，通过堆叠局部运算完成信息的长程传递，但限制了效率的进一步提升。本文提出的非局部UNet新加入了一个全局聚合模块（global aggregation block）来汇聚全局信息得到更加精确的分割结果。最终在婴幼儿脑MR图片中测试了non-local UNet的分割效果。

Section I Introduction

近年来语义分割网络从最初的开山之作FCN走到今天，延伸出了UNet，DeepLab，RefineNet等，其中UNet更是被广泛用于医学图像分割领域。
UNet的基本架构是下采样编码网络+上采样解码网络+中间跃层连接，skip connection很好的将局部和全局的上下文信息结合到了一起
随后基于UNet进行了一系列改进和优化，如加入残差连接提升短程信息的传递；还有将UNet扩展至3D图像的3D-UNet，V-Net以及3D-FCN。
但UNet及其变体存在以下问题：
在编码网络通过一系列卷积核下采样操作堆叠在一起，经输入经过小尺寸的过滤器从而获得局部特征，局部特征级联起来获得长程信息，这些都是基于局部操作。但这严重影响了网络特征提取的效率。因为为了获得全局性特征往往需要搭建深层网络，每次下采样通道数会double一倍从而使得训练参量大大增加，影响效率；而且连续的下采样会导致空间信息的丢失，而这不利于医学图像的精细分割；
另一方面解码网络也是以堆叠的方式进行，将下采样替换为上采样（如转置卷积、unpooling等）用以恢复图像空间信息，缺少空间信息的辅助也不利于图像的复原。
基于上述UNet的两大限制，我们在编码网络部分引入了基于自注意力机制的全局信息聚合模块，不用太深的编码网络就可以完成全局信息的聚合；随后在解码网络中同样引入上采样全局信息聚合模块进行优化。分割实验的结果显示优化后在网络精简的基础上性能获得了进一步提升.

Section II Non-Local UNet

首先介绍基础的UNet框架，然后在不同尺寸、不同下采样/上采样规格的UNet基础上加入全局信息聚合模块组成non-local UNet.
UNet编码网络经过两次下采样用以获取高层次特征（同时降低了空间尺寸）每次下采样后通道数翻倍，最后在bottleneck部分分将信息聚合作为解码网络的输入；相应的在解码网络经过两次上采样恢复空间信息，每次上采样通道数减半。为了帮助解码网络恢复图像空间信息，还会借助skipconnection将同层次对应的编码输出的特征图谱级联（通道concatenate）过来。
但在non-local UNet中是通过summation的方式，这样一是不会增加通道数，减少训练参量，二是可以将这种summation连接看做残差连接集合长程信息，有利于模型训练。

UNet

Residual Block
已经有研究证明了残差连接有利于深层网络的训练，本文这种summation connection的连接方式就可以视作一种长程残差连接；也有研究表明在UNet中加入短程的残差连接可以提升UNet性能；本文除了在下采样中加入残差连接，还在上采样模块中验证了残差连接的有效性。

(a)是标准的残差块（b）是下采样采用的残差块用步长为2的11卷积替换原来的identity mapping（c)是bottom block，可以看到经过了全局聚合模块（d)是上采样采用的残差块，用33,s=3的转置卷积替换identity mapping，以及加入了上采样全局信息汇聚模块。

Global Aggregation Block
卷积这种局部操作操作，某一点的输出只反映其对应感受野的特征；为了完成全局信息的汇聚，某一点的输出应该是输入所有点的反映，比如FC层（但实际应用会导致过拟合，效果不好），而之前有研究证明了self_sttention block可以较好的完成序列的长程依赖关系的捕捉（用于视频分类）。
基于self-attention我们提出global aggregation block将所有特征图谱汇聚出全局信息
操作分3步走：
Step 1:切块 将输入拆成value.key.querymatrix 3部分。首先value和key做dot-production会得到一个标量权重值，随后还会经过softmax。

Step2:注意力机制的实现。将上一步得到的注意力图施加给value matrix。

Step3:维度还原。将施加了注意力的输出按照需要还原为特定维度，作为全局聚合模块的输出。

Section III实验结果讨论

随后将non-local UNet在脑核磁共振图像上进行分割，分割出脑脊液灰质白质（CSF,GM,WM)区域。测试了non-local unet与baseline的性能对比，消融实验验证每一部分的全局信息聚合模块到底是起什么作用，以及在训练和推断阶段如何权衡分割精度和推断时间。
baseline model:CC-3D-FCN,skip-connection是concatenate的方式。
评价指标：Dice ratio,3D-MHD
由于上述两个指标都是进行二分类，所以我们的模型最后会输出4张分割图谱分别显示CSF,GM,WM的分割结果。

Part B训练策略
训练：
在每个全局聚合模块中使用0.5的dropout，图片切成323232大小的patch进行训练（用于增加训练数据，减少对memory的要求），除此之外没有进行其他数据增强。
推断：也是对32大小的patch进行分割，由于滑窗过程中有重叠，因此取得是重叠部分的均值作为最终的分割结果。所以滑窗过程中的重叠步长（overlapping step size)就是一个很重要的超参数，影响着模型的分割精度、推断时间。（步长越小，精度越高但会增加patch数目增加从而带来推断时间的增加），最终step-size取8
Part C 分割结果对比
根据可视化结果，可以看到non-local UNet可以捕获更多细节，而且收敛的更快；在参量方面比baseline减少了28%的同时还获得了性能的提升

Part D 消融实验
Model1:3D-UNet without short-range residual connection
Model2:3D-UNet with short-range residual connection
Model3:上采样网络中将第一个block用2(d)上采样的全局信息聚合模块替换
Model4:上采样网络中全用2(d)上采样的全局信息聚合模块替换
Model5:3D-UNet,bottleneck部分用聚合全局信息模块代替
可以看到聚合模块对模型分割精度的提升（0.5-0.6）

Part E 其他参数的影响
overlapping step size
实验测定了step-size = 4,8,16,32分别对应11880,1920,387,80个patch。
最后选定step-size=8or16可以在性能和时间上取得较好的兼顾。

Patch-size
测试了patch-size = 16,24,32,40,48,最后确定取32.

SectionIV conclusion
本文基于UNet大多考虑局部信息忽略全局信息的限制设计了全局信息聚合模块用来提升分割效果，这一模块的好处是可以方便的设计大小，插入到下采样、上采样中用于聚集全局信息。