OSFormer阅读笔记

前言

OSFormer是华科发表在ECCV 2022的一篇伪装实例分割的文章，提出了一种新的位置感知的一段式（one-stage）Transformer框架，采用了一种高效的位置感知Transformer来捕获全局特征，并动态回归伪装实例的位置和主干。OSFormer还进一步设计了多尺度特征融合，来整合从粗到细的多尺度特征，从而突出伪装实例的边缘，产生全局特征。

大量实验结果表明，OSFormer效果优于其他一直模型。此外，OSFormer只需要大约3000张图像进行训练，并且能快速收敛。因此，可以很容易的扩展到其他训练样本较小的下游视觉任务。

1. 模型的特点

OSFormer主要有以下三个特点：

• 基于Transformer的一段式框架。对比二段式的Fast-RCNN（先寻找前景及调整边界框，再进行分类），OSFormer可以直接生成分类结果。
• 位置感知Transformer（LST），可以动态获取不同位置的实例线索，使得模型在有限数据集的情况下（大约3000张图像），快速收敛。
• 由粗到细的特征融合（CFF）。融合主干网络和LST的多尺度低级和高级特征来获得高分辨率的掩码特征。还嵌入反向额边缘注意力模块（REA）来突出伪装实例的边缘信息。

2. 模型结构

OSFormer主要包含四个基本部分：

• 用于提取特征的CNN backbone
• 位置感知Transformer（LST），利用全局特征和位置引导Query来生成实例感知向量
• 由粗到细的特征融合（CFF），融合多尺度高层和低层特征，产生高分辨率的掩码特征
• 用于预测最终分割掩码的动态伪装实例归一化（DCIN）

2.1 CNN backbone

给定一个输入图像$I\in {\mathbb{R}}^{H\times W\times 3}$，使用ResNet-50作为backbone提取特征。为了减少计算量，将最后三个特征图（$C_3,C_4,C_5$）直接展平并相互连接成具有256个通道的序列$X_m$，将$X_m$作为LST的输入。

将$C_2$特征图输入到CFF模块中，作为底层特征以捕获更多伪装实例线索。

2.2 位置感知Transformer（LST）

Transformer可以通过自注意力模块很好地提取全局特征，但是它需要大量的训练样本和计算成本支持。由于伪装实例分割的数据集有限，所以要在保证性能的同时，使模型更快地收敛。

• LST Encoder

与仅具有单一尺度低分辨率的特征输入的DETR不同，LST encoder的输入是多尺度特征$X_m$。
与Deformable DETR不同，在可变形自注意力层之后，为了更好的捕获局部信息并增强相邻token之间的相关性，将卷积操作引入前馈网络，称为混合前馈网络（BC-FFN）。
首先，根据$C_i$的形状将特征向量恢复到空间维度。然后，使用$3\times 3$卷积来学习归纳偏置。最后，用组归一化（GN）和GELU激活函数来形成前馈网络。再经过一个$3\times 3$卷积，将特征图展平为序列。
与SegFormer中的mix-FFN相比，BC-FFN不包含MLP操作和残差连接。
与CvT在每个stage开始时都加一个卷积token向量不同。OSFormer只在BC-FFN中引入两个卷积层。
假设输入特征表示为$X_b$，BC-FFN的数学表示如下：
$$X_b^{{}^{\prime }}=Con{v}^3(GELU(GN(Con{v}^3(X_b))))$$
其中，$Con{v}^3$为$3\times 3$卷积操作。

整个LST encoder的数学表示如下：
$$X_e=BC-FFN(LN((X_m+P_m)+MDAttn(X_m+P_m)))$$
其中，$P_m$表示位置编码，LN表示层归一化，MDAttn表示多头可变形自注意力。

• 位置引导 Queries

Object Queries在Transformer结构中发挥关键作用，它可以作为Decoder的输入，通过Decoder得到输出向量。但是，DETR收敛缓慢的原因之一就是Object Queries 全部初始化为零。因此，OSFormer提出了位置引导Queries。

受到SOLO的启发，OSFormer首先将恢复的特征图 T3-T5 重新调整大小为$S_i\times S_i\times D,i=1,2,3$。（$T3\in {\mathbb{R}}^{\frac{H}{8}\times \frac{W}{8}\times D},T4\in {\mathbb{R}}^{\frac{H}{16}\times \frac{W}{16}\times D},T5\in {\mathbb{R}}^{\frac{H}{32}\times \frac{W}{32}\times D}$）然后，将调整大小后的特征图划分为$S_i\times S_i$的网格。并将其展平以产生位置引导Query，$Q\in {\mathbb{R}}^{L\times D},L={\sum }_{i=1}^3{S}_i^2$。

这样位置引导Query可以利用不同位置的可学习局部特征来优化初始化，并有效聚合伪装区域中的特征。与DETR中的零初始化，或者Deformable DETR中的随机初始化相比，这种Query生成策略可以提高Transformer Decoder中Query迭代的效率，加快训练的收敛速度。

• LST Decoder

LST Decoder将LST encoder产生的全局特征与位置引导Query进行交互，这对产生实例感知特征向量至关重要。
将空间位置编码$P_s$和$P_m$分别添加到位置引导查询$Q_L$和encoder输出$X_e$中。然后用交叉注意力对二者融合。这里与一般的交叉注意力不同。因为所提出的Query中已经包含可学习的全局特征，这里使用的是没有自注意力的交叉注意力。

与LST encoder相似，在可变形注意操作之后也使用了BC-FFN。假设位置引导Quary为$Q_L$，LST decoder的流程总结为：
$$X_d=BC-FFN(LN((Q_L+P_s)+MDCAttn((Q_L+P_s),(X_e+P_m))))$$
其中，$P_s$表示基于特征网格的位置编码。MDCAttn表示多头可变形交叉注意力操作。$X_d$是实例感知表示的输出向量。最后，将$X_d$送到DPIN模块中以预测掩码。

2.3 粗细特征融合（CFF）

如图所示，将多级特征C2、T3、T4和T5作为级联融合的输入。从输入比例$\frac{1}{32}$的T5开始，进行$3\times 3$卷积、GN和2倍双线性插值上采样。然后，加入更高分辨率的特征（$\frac{1}{16}$比例的T4）。以此类推，最后将特征图C2与$\frac{1}{4}$比例特征图进行融合，再将融合后的特征通过$1\times 1$卷积、GN和RELU运算生成用于分割的特征$F$。注意：每个输入特征在第一次卷积后将通道数从256减少到128，然后在最终输出时增加到256个通道。

考虑到伪装实例的边缘特征更难捕捉，设计了一个包含在CFF中的反向边缘注意（REA）模块，用于在迭代过程中监督边缘特征。受卷积块注意力模块启发，输入特征图同时经过平均池化和最大池化处理，然后将二者拼接起来，经过$7\times 7$卷积和Sigmoid激活，这样特征图中的数值变换到了o与1之间。反转注意力权重，将其与融合特征$F_f$进行元素乘法。最后经过$3\times 3$卷积来预测边缘特征。

假设输入特征为$T_i$，每个REA模块可以表示为下式：
$$F_e=Con{v}^3(F_f\otimes (1-Sigmoid(Con{v}^7([AvgPool(T_i);MaxPool(T_i)]))))$$
其中，$Con{v}^7$代表$7\times 7$卷积，“;”代表在通道维度上进行拼接。
总之，CFF为DCIN提供了一个融合掩码特征F，以最终预测每个伪装实例的掩码。

2.4 动态伪装实例归一化（DCIN）

借鉴风格迁移模型（如AdaptIS等），作者采用动态网络的思路，将动态伪装实例归一化引入到最终掩码的预测中。结构如下图。
DCIN的输入是来自LST decoder的输出$X_d\in {\mathbb{R}}^{S^2\times D}$，将其通过全连接层训练得到位置标签。同时，将$X_d$通过多层感知机（MLP）来获得维度为256的实例感知参数。在训练阶段，通过GT来分配正、负位置标签。使用正位置标签的实例感知参数来生成分割掩码。

在测试阶段，在suppress抑制中使用位置标签的threshold > 0.5来过滤小的参数。随后，将过滤后的位置感知参数通过两个线性层获得放射权重$w\in {\mathbb{R}}^{N\times D}$和偏置$\beta \in {\mathbb{R}}^{N\times 1}$。
最终，$w$和$\beta$与融合掩码特征$F$共同用来预测伪装实例。数学表示为：
$$P=U_{\times 4}(wF+\beta )$$

其中，$P\in {\mathbb{R}}^{H\times W\times N}$是预测掩码。$N$代表预测实例的数量。$U_{\times 4}$是4倍上采样。最后，使用SOLOv2中提出的NMS矩阵得到最终实例。

2.5 损失函数

训练过程中总体损失函数可以表示为下式：
$$L_{total}={\lambda }_{edge}{L}_{edge}+{\lambda }_{loc}{L}_{loc}+{\lambda }_{mask}{L}_{mask}$$

其中，$L_{edge}$代表边缘损失，用以监督CFF中不同层级特征图的边缘。边缘损失可以定义为$L_{edge}={\sum }_{j=1}^J{L}_{dice}^{(j)}$，$J$是CFF中监督边缘特征的总层数。$\lambda$是边缘损失的权重，默认设置为1。
由于伪装实例分割任务中类别个数是未知的，因此我们使用每个位置（$L_{loc}$）存在伪装目标的置信度，而不是通常实例分割中的分类置信度。为了实现分割任务，$L_{loc}$是通过Focal loss（2017年CCF-A论文）进行计算的，$L_{mask}$是通过Dice loss进行计算的。为了对总损失函数$L_{total}$平衡，${\lambda }_{loc}$默认设置为1，${\lambda }_{mask}$默认设置为3。

3. 可视化分析

对CFF模块中每个尺度的输入特征图和最终的CFF输出特征$F$进行了可视化。
图（a）是输入图像重叠真实标签；图（b）到（e）是输入到CFF中的CNN backbone和LST encoder产生的特征图。图（f）是CFF的最终掩码特征输出$F$。

将OSFormer与Mask R-CNN和SOLOv2进行定性比较，如下图。从第一行可以看出，OSFormer可以在多个实例的情况下轻松区分伪装。从第二行可以看出，OSFormer在捕获细长边界方面表现出色，这归因于REA反向边缘注意力模块对边缘特征的增强。

参考

【ECCV 2022】OSFormer论文推介