Non-local U-Nets for Biomedical Image Segmentation

论文地址
Non-local U-Nets
开源代码地址（tensorflow代码）：
github
BCHW = Nx3x512x896，其中
• B表示批次数目
• C表示图像通道
• H表示图像高度
• W表示图像宽度
现有的模型一般基于U-Net，它依赖重复叠加的局部算子来聚合远程信息。这样做会限制模型的训练效率和最终效果。这篇文章提出了非局部的U-Nets网络架构，提出了全局聚合块的应用，它能够融合来自任何大小的特征映射的全局信息。通过在三维多模等强度婴儿脑磁共振图像分割任务上进行试验证明这个模型参数少，计算速度快，并且具有更好的分割效果。
数据集
等强度婴儿脑磁共振图像分割任务

德克萨斯A&M大学(Texas A&M University)，简称TAMU。该校建立于1876年，位于美国德克萨斯州的卡城是旗舰主校区，该校在Galveston（Texas）、Corpus Christi（Texas）、Commerce（Texas）、Kingsville（Texas）、San-Antonio（Texas）和 Doha（Qatar）等还有分校。TAMU是公认世界百强名校，

北卡罗来纳大学教堂山分校（University of North Carolina at Chapel Hill），简称UNC，创建于1789年，是世界顶级研究型大学，美国历史上第一所公立大学。该校是美国最顶尖的五所公立大学之一，被誉为“公立常春藤”和“新常春藤”

u-net存在的两个缺陷
1． 图像分割要求对每个像素进行准确地预测，空间信息的损失会导致分割图不准确。
unet编码器均为局部操作，无法整合全局信息，同时下采样丢失空间信息，尤其对生物医学图像不利。unet编码器通常会堆叠卷积层与下采样交叉在一起，逐渐减小特征图的空间尺寸。卷积操作、下采样操作都是局部操作，运用小卷积核进行特征提取。通过级联的方式叠加卷积和下采样操作产生较大卷积核，因此能够聚集较大范围信息。由于生物医学图像分割通常受益于广泛的上下文信息，因此大多数模型需要深层的编码器，即堆叠更多的局部操作。这样会引入大量训练参数，尤其是在需要更多下采样的时候，因为通常这样特征映射的通道数会加倍此外，下采样会丢失更多的空间信息，这些信息对于生物医学图像分割至关重要
2.decoder的形式与encoder部分正好相反，包括若干个上采样运算，使用反卷积或插值方法，他们也都是local的方法。
unet解码器同样为局部操作，上采样很难恢复全局信息，尤其是无法从编码器有效获取全局信息。unet解码器上采样如反卷积，反池化操作都是局部操作，上采样恢复细节信息需要全局信息。如果不考虑全局信息就很难做到这一点。
上采样的过程涉及恢复图像的空间信息，如果只是局部信息而不考虑全局信息就很难做到。

比较意义
相比U-Net，卷积层数减少，图像下采样倍率从16倍变成4倍，保留了更多空间信息。

框架图
网络使用的3d结构，上图示例输入为两个通道，输出为4类，每一次下采样通道加倍，反之上采样通道减半，输入首先经过一个编码输入块，它提取低层特征。接着利用两个下采样块来减少空间尺寸，以此来获得高阶特征。在这之后，底层块聚合全局信息并产生编码器的输出。对应的，解码器使用两个上采样块来恢复分割输出的空间大小。
短连接采用相加操作。优点有两点：
3.1.跳跃连接不会增加特征图的数量，从而减少了参数
3.2.相加可以视作残差连接，在训练中任然有效
encoder和decoder之间的skip connections用相加的方式，不是拼接的方式，让模型推理速度更快。作者将原始u-net的concatenation -> sum的原因：1求和不会增加特征图的数量，因此会减少下一层中可训练参数的数量
。2具有求和的skip connection可以被视为远程残差连接，可以促进模型的训练。

残差快
提出的网络结构中相加操作等价于远程残差连接模块，并且基于global aggregation block，作者提出了四种残差网络，都采用了pre-activation模式。
（a）常规残差块。 该块用作输入块。 输出块由该块+步长为1的1x1x1的卷积构成。此外，在跳跃连接求和之后，我们插入一个这样的块。
（b）是下采样残差块。 将此块用作下采样块。
（c）说明了我们的底部模块。 残差连接应用了我们提出的全局聚合块。
（d）上采样残差块,（b）中的下采样块相似，恒等残差连接被步长为2的3×3×3反卷积代替，另一个分支是上采样全局聚合块。 使用此块作为上采样块

全局聚合块
实现全局信息融合：输出特征图的每个位置都应取决于输入特征图的所有位置。 （与卷积和反卷积等局部操作相反）,可以通过self-attention来聚合(aggregation)图像特征图的全局信息。Attention函数的本质可以被描述为一个查询（query）到一系列（键key-值value）对的映射
在计算attention时主要分为三步
x，y：输入输出
Conv_1N：输出通道为N的点卷积
Unfold(·)：将D × H × W × C展开为(D ×H ×W)×C张量
QueryTransformCK (·)： 生成CK feature maps的任何操作
CK, CV：是表示键和值的维度的超参数
K,V: 假设X的尺寸为D×H×W×C，则K和V的尺寸分别为（D×H×W）×CK和（D×H×W）×CV

代码里在Unfold之前有Multi-Head（Attention Is All You Need）的操作，不过在论文中没有说明，实际上是把通道分为N等份。Unfold是把Batch,height,width,N通道合并,Q（BHqWqNck），K（BHWNck），V(BHWN*cv）。

3.接下来是经典的点积attention操作，得到一个权值矩阵A(A的维度为(DQ × HQ × WQ) × (D × H × W)。O的维度为(DQ × HQ × WQ) × CV)，用于self-attention的信息加权，分母Ck是通道数，作用是调节矩阵的数值不要过大，使训练更稳定（这个也是Attention Is All You Need提出的）。最后权值矩阵A和V点乘，得到最终的结果（(BHqWqN)*cv），可见输出的height和width由Q决定ÿ