Transformer meets UNet
简单来说,就是分为两个流,用CNN来获取局部信息,用tranformer获取全局上下文信息,Transformer模块生成一个图像级上下文表示(ICR)来构建图像级的空间依赖图,并生成区域重要性系数(RIC)来模拟每个区域的重要性。
模型结构如下:
在第一个流当中,我们应用CNN来获取图像(但好像是个二维的图像),经过半U-Net的模型之后
边界热图的获得:
其中σ表示sigmoid激活,而Convb(.)表示1 × 1核卷积操作
2、远程信息的表示:
上面说到,要通过transfomer来实现远程上下文的依赖表示,所以将图像切割成了p*p*c的patch(注意c是通道数,而不是深度,因为他是个二维的),每个patch变成了flattened并且无重叠部分,故每部分的序列长度是
经过Patch Encoder将这些补丁嵌入到K维的映射当中,为了保持每个patch的空间信息,我们学习到一个一维位置嵌入Ipos∈RN ×K,随后加入到patch嵌入中以保持位置信息
其中I∈R(p2·C)×K表示投影的patch嵌入。然后,我们利用包含多头自注意(MSA)和多层感知器(MLP)的Transformer块堆栈来学习远程上下文表示。MSA层由M个并行的自注意头组成,用于缩放嵌入的补丁:
然后,MLP学习远程依赖通过,其中Norm()表示层归一化[5】,
表示d维空间中的编码语义表示。除了编码特征之外,我们还通过重塑(Re)特征并应用1 × 1卷积操作(Conv1)来建模图像级上下文表示(ICR):
我们使用ICR在图像级构建空间依赖图,然后对CNN模块生成的特征集进行归一化。我们进一步定义区域重要系数(RIC)来模拟前景像素在每个区域的分布。RIC系数的目的是提供一个监督信号,指导上下文注意模块确定重要区域(Eq. 3)。ConvR显示了一个1 × 1的卷积操作。
3、环境注意模块
为了自适应的为了自适应地聚合提取的特征,我们提出了上下文注意模块。每个特性集的重要性应该与任务一致,我们提出的模块利用以下两级规范化步骤:首先,它重新校准CNN表示以进行像素级对象理解,然后执行空间规范化以选择性地强调特征集中的远程上下文依赖关系。在挤压和激励之后[19],我们将通道标准化权重(ch)定义为:
其中GAP表示应用于CNN特征(f)的全局平均池化操作,W1和W2是学习参数,δ和σ是Sigmoid和ReLU激活函数。我们形成归一化特征:
为了强调边界区域,我们将边界表示添加到归一化特征中,,边界特征的目的是强调边界区域,引导模型精确分离重叠对象(如对象级交互)。接下来,使用从Transformer模块派生的特征集,我们执行空间归一化。为此,首先我们将RIC系数与相应的区域相乘,根据区域重要性
对表示进行缩放。为了进一步整合远程依赖关系,我们将ICR表示与fsn连接起来,然后应用卷积核,然后是批处理归一化(BN)和激活函数来执行非线性聚合:
Conv是1×1卷积运算。结果特征集包含局部语义和全局上下文表示,它们被选择性地组合起来执行语义分割任务。随后,我们将解码器块应用于提取的特征来预测分割掩码,
。我们在训练过程中优化的联合目标损失函数为:
其中,Lsegmentation计算的是预测值和真实值之间的CELoss,Lboundary计算的是边界的二值预测的CELoss,LRIC计算的是每个图像patch中前景像素分布与相应预测像素分布之间的MSE损失