背景
深度学习在各种生物医学图像分割任务重显示出巨大的应用前景。现有的模型一般基于U-Net,它依赖重复叠加的局部算子来聚合远程信息。这样做会限制模型的训练效率和最终效果。这篇文章提出了非局部的U-Nets网络架构,提出了全局聚合块的应用,它能够融合来自任何大小的特征映射的全局信息。通过在三维多模等强度婴儿脑磁共振图像分割任务上进行试验证明这个模型参数少,计算速度快,并且具有更好的分割效果。
介绍
U-Net的缺点
U-Net是由一个下采样编码器和一个上采样解码器以及它们之间的跳过连接组成的。它通过编解码过程来
U-Net主要有两个局限性:
- 编码器通常叠加卷积和池化操作以逐步减小特征图的大小。这种方式会引入大量参数,降低模型的效率。此外,下采样会不断丢失空间信息,也会影响最终的分割效果。
- 上采样过程涉及到空间信息的恢复,如果不考虑全局信息很难实现。
创新点
文章中针对U-Net的这些缺点,进行了几点创新。
- 提出了一个基于自我主义算子的全局聚合块,它能够在没有深入编码器结构的情况下聚合全局的信息。
- 将这个全局聚合块应用到上采样过程中,也一定程度上解决了上述第二个问题。
Non-local U-Nets
基于U-Net框架的示意图如下:
输入首先经过一个编码输入块,它提取低层特征。接着利用两个下采样块来减少空间尺寸,以此来获得高阶特征。在这之后,底层块聚合全局信息并产生编码器的输出。对应的,解码器使用两个上采样块来恢复分割输出的空间大小。
文章与U-Net的一点不同之处是在跳过连接这里,它在跳过连接之后使用求和方法去结合两个特征图的信息,而不是U-Net中使用的串联。使用求和的好处是:首先求和不会增加特征图的数量从而减少参数数量。其次带有求和的跳跃连接可以被认为是远程剩余连接,它能够促进模型的训练。
模型中使用四个不同的残差块来形成一个完整的残差网络。
- a是一个规则的残差块,它有两个连续的卷积层,每个卷积层之前使用ReLU激活函数进行batch normalization。这个块在框架结构中作为输入块。输出块由这个块和步长为1的1x1x1卷积组成。此外在跳过连接之后,也会插入一个这样的块。
- b是下采样残差块。利用步长为2的1x1x1卷积代替一致剩余连接从而调整特征图的大小。
- c是底部模块,在提出的全局聚合块上应用一个残差连接。
- d是上采样块。与b类似,将一致残差连接替换为步长为2的反卷积,另一个分支是上采样全局聚合块。
全局聚合块
全连接层也有全局属性,但是其容易过拟合。收到Vaswani等人启发,指出图像特征图的全局信息可以通过自注意块进行聚合。
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
X为全局聚合块输入,Y为输出。
- 该结构中,输入图为X(BxHxWxC),经过QueryTransform和1x1卷积,转换为Q(BxHqxWqxCk),K(BxHxWxCk),V(BxHxWxCv)。QueryTransform可以为卷积,反卷积,差值等方法,最后输出结果的H与W将与这个值一致。
- 代码里在unfold之前有Multi-Head的操作,不过在论文中没有说明,实际上是把通道分为N等份。Unfold是把batch,height,width,N通道合并。(生成query (Q), key (K) and value (V ) 矩阵。Unfold之前有Multi-Head操作,Multi-Head将通道拆分,Q的维度为(DQ × HQ × WQ) × CK)
- 接下来是经典的点积attention操作,得到一个权值矩阵A((BxHqxWqxN)*(BxHxWxN)),用于self-attention的信息加权,分母Ck是通道数,作用是调节矩阵的数值不要过大,使训练更稳定。最后权值矩阵A和V点乘,得到最终的结果((BxHqxWqxN)xCv),可见输出的height和width由Q决定,通道数由V决定。(自注意机制。A为注意力矩阵,用Q点乘K的转置,除以根号CK,再使用softmax,A的维度为(DQ × HQ × WQ) × (D × H × W)。O的维度为(DQ × HQ × WQ) × CV。大家可以画画图理解一下2d卷积的情况,本质还是non-local操作)
- 现在很多上采样相关的论文关注于在上采样时用CNN扩大感受野,增加图像局部信息。这篇文章提出的全局聚合块是一个将注意力机制和上/下采样相结合的方法,关注全图信息,感受野更大,可以在其他任务上试用一下,效果如何还是要看实践的结果。
实验结果
作者以三维态等强度婴儿脑磁共振图像分割任务为实验对象,对non-local U-Nets惊醒了实验。该任务是将磁共振图像自动分割成脑脊液、灰质和白质区域。
baseline:CC-3D-FCN(已证明优于传统机器学习方法)
评价指标:
- DR:较大的DR表示更精确的分割。
- MHD:两个物体之间相似度的参数。
表1显示,在脑脊液、灰质、白质的分割效果上,文章提出的方法都要优于基线方法。
表2显示,在脑脊液、灰质、白质的相似度方面,文中提出的方法MHD值较小,因此分割效果更好。
图4给出了分割效果,通过比较红色圆圈中的区域,可以看出文章中提出的方法相比基线模型可以捕捉更多的细节。
图5分别显示了文章中模型与基线模型在运行时的训练和验证曲线,可以发现文章中模型收敛更快,训练损失更低。
表4显示文章中模型的参数相比基线模型的参数数量降低了28%,因此有更高的效率。
表5说明文章中模型的测试时间相较基线模型更短。
消融研究
- Model1:无短程残差连接
- Model2:有短程残差连接,使用了残差块a和b,底层块和上采样块与Model1相同
- Model3:将Model2中的第一个上采样块替换为残差块d
- Model4:将Model2中的两个上采样块全换成残差块d
- Model5:将Model2中的底层块换成残差块c
结果表明,在non-local U-nets中,不同的全局聚合块如何提高性能。
小的重叠步长通常会取得较好的分割效果,但是过小也会影响测试速度。通过实验找到一个较好的重叠步长值,发现8和16是比较好的选择。为4时,patch数量太多,效率低;为32时,DR较小,效果不好。
patch的大小影响不同训练样本的总数。为了选择较好的patch值,进行了相关实验。在5种不同的尺寸中,32取得最好的结果。
结论
提出了一个全局聚合块,可以灵活的用于U-Net中的上采样和下采样过程。在三维多模等强度婴儿大脑磁共振图像分割任务实验表明,文章提出的方法参数更少、运算速度更快、性能也优于先前的模型。
来源:CSDN
作者:strawqqhat
链接:https://blog.csdn.net/strawqqhat/article/details/104835713