笔记03----NeurIPS2024 涨点!SSA:用于语义分割的语义和空间自适应像素级分类器(即插即用)

前言

文章标题：《SSA-Seg: Semantic and Spatial Adaptive Pixel-level Classiffer for Semantic Segmentation》
助力语义分割涨点!SSA:一种新颖的语义和空间自适应分类器，显著提高了基线模型的分割性能，比如SegNeXt、OCRNet和UperNet等模型，而计算成本的增加却很小，代码已开源!
单位:华为诺亚，浙江大学
问题：

1. 语义域特征偏差：由于测试图像的内容和分布可能与训练数据存在很大差异（如背景复杂、对象分布多样），测试图像中的像素特征可能与训练时学习的固定语义原型不一致，导致分类性能下降。
   举一个例子：金斧头，银斧头，你却只认得铁斧头

3. 空间域信息丢失：传统方法未充分利用目标对象的空间结构信息，导致边界模糊以及对小目标的识别能力不足。

创新：
语义和空间自适应分类器（SSA-Seg）：

• 空间原型自适应（SPPA）：利用位置编码获得空间特征，通过粗分割掩码计算空间域中心，建模像素特征与空间之间的关系，生成空间原型，改进像素与空间原型之间的关系，提升边界和小目标分割性能。
• 语义原型自适应（SEPA）：使用粗分割掩码计算语义域中心，生成自适应的语义原型，改进像素与语义原型之间的相似性，提升语义类别对齐能力。

源码下载

github：https://github.com/xwmaxwma/SSA-Seg

实现

SPPA：
对于从骨干和分类头输出的语义特征$S_f$，我们首先通过位置编码生成空间特征$P_f$。然后，我们保留原始的1x1卷积以生成粗级Mask:$M_c$。在$M_c$的引导下，空间域的中心$P_c$，并将其与一个可学习的位置编码$P_p$融合生成空间原型
SEPA：
对于从骨干和分类头输出的语义特征$S_f$，我们保留原始的1x1卷积以生成粗级Mask:$M_c$。在$M_c$的引导下，语义中心$S_c$，并将其与固定的语义原型$S$和$S_c$融合，以生成语义原型$S_p$
SSA：
（空间原型$P_p$+语义原型$S_p$）*（语义特征$S_f$+空间特征$P_f$）=Output

class SEPA_SPPA(nn.Module):
    def __init__(self, num_classes, channels):
        super(SEPA_SPPA, self).__init__()
        self.channels = channels  # 通道数，用于特征图的维度
        self.num_classes = num_classes  # 类别数，用于位置嵌入和预测

        # SEPA模块：生成特征位置编码
        self.get_feat_pos = nn.Conv2d(self.channels, self.channels, 3, 1, 1, bias=True, groups=self.channels)  # 使用深度可分离卷积生成特征位置编码
        self.center_pos = nn.Embedding(self.num_classes, self.channels)  # 可学习的向量类别位置嵌入，大小为(类别数, 通道数)
        self.center_pos_proj = nn.Sequential(
            nn.Linear(self.channels * 2, self.channels // 2, bias=False),  # 降维
            nn.ReLU(inplace=True),  # 激活函数
            nn.Linear(self.channels // 2, self.channels),  # 恢复到原通道维度
        )
        self.feat_proj = nn.Identity()  # 特征映射，不做额外操作

        # SPPA模块：生成语义原型
        self.center_content_proj = nn.Sequential(
            nn.Linear(self.channels * 2, self.channels // 2, bias=False),  # 原型降维
            nn.ReLU(inplace=True),  # 激活函数
            nn.Linear(self.channels // 2, self.channels),  # 恢复到原通道数
        )
        self.center_proj = nn.Linear(
            self.channels, self.channels, bias=False
        )  # 最终类中心特征的线性映射

    # 获取中心位置编码
    def get_center_pos(self, attn, feat_pos):
        """
        根据注意力图(attn)和特征位置编码(feat_pos)计算中心位置编码
        """
        center_pos = self.center_pos.weight  # 获取类中心位置嵌入
        b, k, h, w = attn.size()  # 批量大小、类别数、高度、宽度
        c = feat_pos.shape[1]  # 特征图通道数

        # 将注意力和特征图展平，用于矩阵计算
        attn = attn.reshape(b, k, -1)  # (b, k, h*w)
        feat_pos = feat_pos.reshape(b, c, -1).permute(0, 2, 1)  # (b, h*w, c)

        # 归一化注意力并计算中心位置
        attn = F.softmax(attn, dim=-1)  # 在像素维度上进行softmax
        center_pos = center_pos.unsqueeze(0).repeat(b, 1, 1)  # 扩展为(b, k, c)
        center_pos = torch.cat(
            [center_pos, torch.matmul(attn, feat_pos)], dim=-1
        )  # 拼接原始位置编码和计算出的编码 (b, k, 2c)
        center_pos = self.center_pos_proj(center_pos)  # 投影到中心位置编码 (b, k, c)
        return center_pos

    # 生成像素分类预测
    def get_pred(self, x, proto):
        """
        根据特征图和语义原型生成像素分类预测
        """
        b, c, h, w = x.size()  # 批量大小、通道数、高度、宽度

        if len(proto.shape) == 3:  # 如果原型是三维的 (b, k, c)
            cls_num = proto.size(1)  # 类别数
            x = x / (torch.norm(x, 2, 1, True) + 1e-12)  # 对特征进行L2归一化
            proto = proto / (torch.norm(proto, 2, -1, True) + 1e-12)  # 对原型进行L2归一化
            x = x.contiguous().view(b, c, h * w)  # 展平特征 (b, c, h*w)
            pred = proto @ x  # 点乘生成预测 (b, cls, h*w)
        elif len(proto.shape) == 2:  # 如果原型是二维的 (k, c)
            cls_num = proto.size(0)  # 类别数
            x = x / (torch.norm(x, 2, 1, True) + 1e-12)  # 特征归一化
            proto = proto / (torch.norm(proto, 2, 1, True) + 1e-12)  # 原型归一化
            x = x.contiguous().view(b, c, h * w)  # 展平特征
            proto = proto.unsqueeze(0)  # 扩展维度为 (1, cls, c)
            pred = proto @ x  # 点乘生成预测 (b, cls, h*w)

        pred = pred.contiguous().view(b, cls_num, h, w)  # 恢复为原始大小
        return pred * 15  # 放大预测值

    # 前向传播
    def forward(self, feat, pred, proto):
        """
        输入特征图、注意力图和原型，输出像素分类预测
        语义特征feat
        pred是预测输出值
        proto固定的语义原型，通过调用主干网络的分割头权重赋值self.conv_seg.weight.squeeze()
        """
        # SPPA模块
        feat_pos = self.get_feat_pos(feat)  # 获取特征位置编码
        center_pos = self.get_center_pos(attn=pred, feat_pos=feat_pos)  # 空间原型
        # SEPA模块
        raw_x = feat.clone()
        b, c, h, w = raw_x.shape[:]
        pred = pred.view(b, proto.shape[0], h * w)
        pred = F.softmax(pred, 1)  # b, n, hw
        pred_proto = (pred @ raw_x.view(b, c, h * w).permute(0, 2, 1)) / (pred.sum(-1).unsqueeze(-1) + 1e-12)

        pred_proto = torch.cat([pred_proto, proto.unsqueeze(0).repeat(pred_proto.shape[0], 1, 1)], -1)  # b, n, 2c
        pred_proto = self.center_content_proj(pred_proto) #语义原型
        # 语义特征feat，feat_pos空间特征
        feat = self.feat_proj(feat+feat_pos)  # 对特征图进行投影
        # pred_proto语义原型 center_pos空间原型
        center = self.center_proj(pred_proto + center_pos)  # 结合语义原型和位置编码生成类中心特征
        pred = self.get_pred(feat, center)  # 根据类中心特征生成像素分类预测
        return pred
if __name__ == "__main__":
    # 初始化模型
    num_classes = 21  # 类别数
    channels = 256  # 特征图通道数
    model = SEPA_SPPA(num_classes=num_classes, channels=channels)  # 实例化模型

    # 构造测试输入
    b, c, h, w = 4, 256, 64, 64  # 批量大小、通道数、高度、宽度
    k = num_classes  # 类别数
    feat = torch.randn(b, c, h, w)  # 模拟特征图
    attn = torch.randn(b, k, h, w)  # 模拟注意力图
    proto = torch.randn(k, c)  # 模拟语义原型

    # 前向传播测试
    pred = model(feat, attn, proto)  # 获取预测结果
    print(f"预测结果形状: {pred.shape}")  # 打印预测结果的形状

实验

SSA-Seg在最新通用(顶部)和轻量化(底部)方法上的性能比较。FLOPS(G)的计算基于ADE20K和COCO-Stuff-10K的512x512输入大小，以及PASCAL-Context的480x480输入大小。延迟(ms)的计算基于V100 GPU的512x512的输入大小。绿色数字表示相对于基线的增加