YOSO:实时全景分割网络

算法小乔

已于 2023-07-18 10:35:57 修改

阅读量1.2k

点赞数 4

文章标签：深度学习人工智能

于 2023-04-12 09:02:45 首次发布

本文链接：https://blog.csdn.net/weixin_40723264/article/details/130098335

版权

（1）插值优先聚合 (Interpolation-First Aggregation，IFA)

（2）卷积优先聚合(Convolution-First Aggregation, CFA)

一、概要

本文提出YOSO,一个实时的全景分割框架。YOSO通过全景Kernel和图像特征图之间的动态卷积进行分割预测，该方法处理实例和语义分割任务时，只需要分割一次。

为了减少计算开销，设计了一个用于特征图提取的特征金字塔聚合器，以及一个用于全景内核生成的可分离动态解码器。

其中：

聚合器：以卷积优先的方式重新参数化插值优先模块，这显著加快了Pipeline的速度，而没有任何额外的成本。

解码器：通过可分离的动态卷积执行多头交叉注意力，以获得更好的效率和准确性。

二、网络结构

1. 整体结构

如上图所示，YOSO整体上是有一个特征金字塔和一个可分离动态解码器组成。

其Backbone采用的是ResNet，从输入图像中提取多级特征图，特征金字塔FPN将多级特征图压缩并聚合为一个特征图。

然后，通过可分离动态解码器生成全景kernel，进而进行mask的预测和分类。

2. 特征金字塔

如上图所示，C2，C3，C4，C5是backbone的多级特征图，通过FPN + DCN（形变特征金字塔）的方式，对多级特征进行增强和融合，具体操作如下：

对于C2到C5，首先采用1x1的卷积对多级特征图进行通道压缩，然后C3到C5分别进行DCN和upsample操作，从上到下的特征融合，分别得到P5、P4、P3、P2.

得到 P2、P3、P4、P5之后进行特征聚合,得到同P2同shape的特征图S

YOSO指出了两种方法，分别是 插值优先的IFA 和 卷积优先的 CFA。下面分别进行介绍：

（1）插值优先聚合 (Interpolation-First Aggregation，IFA)

首先P3到P5，分别上采样到P2大小，然后Concat通道合并，最后通过1x1卷积进行通道间的特征融合。

（2）卷积优先聚合(Convolution-First Aggregation, CFA)

首先P2到P5分别进行1x1卷积，然后P3到P5卷积后的结果分别上采样到P2大小，最后直接进行Add操作。

两种方法的1x1卷积中，bias=False。

总结：上述两种方式，最终的效果是相当的；但是CFA的计算量更低一些。最终采用了CFA.

import torch.nn.functional as F 
import torch.nn as nn 
# 经过FPN结构得到x2,x3,x4,x5
class CFA_IFA(nn.Module):
    def __init__(self,in_channels_list,out_channels,mode='cfa'):
        super(CFA_IFA,self).__init__()
    	self.mode = mode 
        if self.mode=='cfa':
            self.conv_a5 = nn.Conv2d(in_channels_list[-1],out_channels,1,1,0)
            self.conv_a4 = nn.Conv2d(in_channels_list[-2],out_channels,1,1,0)
            self.conv_a3 = nn.Conv2d(in_channels_list[-3],out_channels,1,1,0)
            self.conv_a2 = nn.Conv2d(in_channels_list[-4],out_channels,1,1,0)
        else:
            self.fuse_conv = nn.Conv2d(4*out_channels,out_channels,1,1,0)
        self.output_conv = nn.Conv2d(out_channels,out_channels,3,1,1)
    
    def forward(self,x2,x3,x4,x5):
        if self.mode=='cfa':
            # CFA
            x5 = F.interpolate(self.conv_a5(x5), scale_factor=8, align_corners=False, mode='bilinear')
            x4 = F.interpolate(self.conv_a4(x4), scale_factor=4, align_corners=False, mode='bilinear')
            x3 = F.interpolate(self.conv_a3(x3), scale_factor=2, align_corners=False, mode='bilinear')
            x2 = self.conv_a2(x2)
            x = x5 + x4 + x3 + x2 + self.bias      
            x = self.output_conv(x)
            return x
		else:
            # IFA
            x5 = F.interpolate(x5, scale_factor=8, align_corners=False, mode='bilinear')
            x4 = F.interpolate(x4, scale_factor=4, align_corners=False, mode='bilinear')
            x3 = F.interpolate(x3, scale_factor=2, align_corners=False, mode='bilinear')
            x = torch.concat([x5,x4,x3,x2], dim=1)
            x = self.fuse_conv(x)
            x = self.output_conv(x)
			return x

3. 可分离动态解码器

为了生成精确的分割Kernel，之前的方法通常采用密集的预测器(如sigmoid)或重型的Transformer解码器。

YOSO，则采用了一种轻量可分离的动态解码器来生成kernel，在保持高精度的同时加快了kernel的预测与生成。结构如下：

如图所示，由三部分组成：Pre-Attention(预注意力模块)、可分离动态卷积注意力模块、Post-Attention(后注意力模块)

(1) Pre-Attention模块

如上图的红色框所示:

a. 作用：

从主干网络的FPN的聚合特征S(bxc x h x w)选择性地(自适应地)提取关键信息。

b. 操作流程：

<1> 获取映射矩阵Mask

S的shape为b x c x h x w, 即通道数目为c, 高为h，宽为w 。

首先通过Conv2d卷积对S进行处理，kernel_size=1,调整通道数目为 n ( n 指的是proposal kernels的数目，即每张图像设定一个最大的候选结果数目proposal kernels number，后期在通过处理从候选结果中筛选 )得到卷积结果 A (shape为b x n x h x w)

然后对卷积结果 A 进行Sigmoid处理，并设定阈值0.5，作为一个注意力矩阵Mask shape为b x n x h x w（有点空间注意力的意思）

最后通过Msak对S进行注意力映射。

<2> 特征映射获得self-Attention的V

映射流程：

Mask shape: [b x n x h x w] -> reshape [b , n , hw] , 表示 n 个 proposal kernels , 每个 proposal kernel 有 hw 个像素点，将每个像素点映射到对应的n个proposal kernels。

S shape : [b x c x h x w] -> reshape [b,c,hw] ，表示 hw个像素点，每个像素点有c个通道数。

V 映射结果，shape : [b x n x c], 相当于对每个 kernel 进行特征选择，对于每个 kernel 从 特征矩阵 S 中选择 c 个特征值，对应于上图中的 Masked Features

公式如下所示：

上图中的Proposals Kernels 指的就是2D卷积的卷积核Q, 估计当时作者这样绘图的初心是为了更形象，因为 Q 在后面的stage中也会用到！！！。

注意！！ 公式中的Q 也就是后面动态卷积注意力(Dynamic Convolution Attention )Decoder中的Q。它来自于上述2D卷积的卷积核 shape [n,c,1,1] -> reshape [n,c], 与V的shape一致(不考虑batch)。

对应回图像，如下所示：

可见Q既应用于卷积操作，也应用于下面的注意力模块。

代码如下：

import torch 
import torch.nn as nn 
# 首先声明一个卷积函数，用于上面公式中的红框操作
class YOSOHead(nn.Module):
    def __init__(self,in_channels,num_proposals):
        super(YOSOHead,self).__init__()
        
        '''
        in_channels: 主干网络输出S的通道数
        num_proposals: 设置的proposals数目，kernel的个数
        '''
        self.num_proposals = num_proposals
        # 声明一个Conv用于卷积操作，同时会获取卷积权重，用于后面注意力模块的Q
        self.kernels = nn.Conv2d(in_channels,num_proposals,kernel_size=1)
        
    def forward(self,features,xxx):
        B,C,H,W = features.shape
        # features : b x c x h x w 
        # mask_preds: b x n x h x w 
        mask_preds = self.kernels(features)# 卷积操作
        # sigmoid处理并与阈值比较获得mask 
        # b x n x h x w 
        soft_sigmoid_mask = mask_preds.sigmoid()
        nonzero_inds = soft_sigmoid_mask > 0.5
        hard_sigmoid_mask = nonzero_inds.float()
        # V = r(A)r(S).T
        # b x n x c 
        V = torch.einsum('bnhw,bchw->bnc', hard_sigmoid_masks, features)
        # ------------------ 以上 Pre-Attentinon ------------------ # 
        
        # 获取注意力的Q
        # n x c x k x k   k = 1
        proposal_kernels = self.kernels.weight.clone()
        # reshape  b x n x c x 1 x 1 -> b x n x (cxkxk)= b x n x c  与V的shape相同
        Q = proposal_kernels[None].expand(B,*proposal_kernels.size()).view().view(B,self.num_proposals,-1)

（2）可分离动态卷积注意力模块

<1>传统交叉注意力模块

如下图所示：

上图为传统的多头交叉注意力模块，将输入 Q和V sape[n,c] ,拆分为t个head，进行交叉注意力计算，公式如下：

其中，Wo ，shape[c,c] 是一个线性映射矩阵，每个Head内部分注意力定义计算如下：

其中shape均为 [c,c/t], 是第i个Head的投影映射矩阵。结果中的Ki 表示相关矩阵，即Attention矩阵shape [b,n,n]，Vi 表示第i 个Head的输入特征,shape [b , n , c/t], Hi表示第i个Head的输出，shape为[b,n,c/t].

传统方法的不足：

虽然可以提升性能，但是计算量很大。

<2> 一维卷积进行多头交叉注意力

从上面叙述可知，多头交叉注意力主要涉及三个基本操作：多头映射 (如下图红色框)、 注意力计算(如下图绿色框)、 注意力融合(如下图蓝色框)

其中表示将 Q 或 V 从维度 c 映射到维度为 c/t 的t个不同的Head中。而KiVi 表示第i 个Head中，特征矩阵在不同proposal kernels中的交叉融合。

本文方法中，QV和KV的操作，我们采用1D 卷积代替，进行多头交叉注意力，使模块更为轻量化，如下图所示：

Q来自于上面Per-Attention中的Proposal Kernels 其shape为[n,c,1,1] -> expand[b,n,c,1,1] -> reshape [b,n,c] ,V 为 Per-Attention中得到的映射矩阵其shape为[b,n,c]。

如上图，Q分别进行两个全连接得后到两个结构Q' （作为第一个卷积的卷积核）和 Q'' (作为第二个卷积的卷积核)。

因为两个都是对Q的全连接

a. 由 Q [b,n,c] 经过 nn.Linear(c,3)即QWd -> shape[b,n,3](如上图红框)

Q' = depth_conv_weight shape为[b,n,3]

reshape-> [b,n,1,3],通过一维深度可分离卷积，通道卷积部分，作为第一个卷积的卷积核,与V [b,n,c] ->reshape [b,1,n,c], 进行1D卷积处理，得到O' [b,n,c]；

b. 由 Q [b,n,c] 经过 nn.Linear(c,n)即QWp -> shape[b,n,n](如上图黄框)

Q'' = point_conv_weight shape 为[b,n,n]

reshape-> [b,n,n,1],后面通过一维深度可分离卷积，融合部分，作为卷积核，处理O' [b,n,c] -> reshape [b,1,n,c]最终得到O'' [b,n,c]；

卷积操作时，将Batch内每张图像单独处理，这样 Q' [n,1,3]

对V中Bacth每个样本，单独进行一维卷积处理(如上图绿框)，且分组为groups=n,模仿DepthWise Convolution的3x3单通道卷积部分，得到Decoder 的Attention矩阵Attention shape为[1,n,c],公式如下所示：

其作用与下图中的绿色框的效果相当。

Q'' [n,n,1] , 其作为卷积核，对Attention中每个样本进行一维卷积处理，模仿DepthWise Convolution的1x1的通道间信息融合部分，公式如下所示：

运算过程中，O' 为第一个conv的输出，shape为[1,n,c] , O''作为第二个conv的输出, shape为[1,n,c],具体操作看后面的代码

其作用与下图公式的红色框作用相当。

batch内所有图像的结果合并，并进行Norm正则化处理，得到输出 H shape[b,n,c]

最后的输出结果再与输入进行融合,得到输出结果O shape[b,n,c]如下图箭头所示：

从上图可知，该部分可以由N个Blocks组成，在论文方法中N=2.以上介绍得只是第一个block，从第二个开始，输出的数据会有所变化，需要注意。

import torch 
import torch.nn as nn 
# 首先声明一个卷积函数，用于上面公式中的红框操作
class YOSOHead(nn.Module):
    def __init__(self,in_channels,num_proposals):
        super(YOSOHead,self).__init__()
        
        '''
        in_channels: 主干网络输出S的通道数
        num_proposals: 设置的proposals数目，kernel的个数
        '''
        self.num_proposals = num_proposals
        self.in_channels = in_channels
        self.conv_kernel_size_2d = 1
        # 声明一个Conv用于卷积操作，同时会获取卷积权重，用于后面注意力模块的Q
        self.kernels = nn.Conv2d(in_channels,num_proposals,kernel_size=self.conv_kernel_size_2d)
        # 第一个block
        self.f_atten = DySepConvAtten(in_channels)
        self.f_dropout = nn.Dropout(0.0)
        self.f_atten_norm = nn.LayerNorm(in_channels)
        # 第二个block
        self.k_atten = DySepConvAtten(in_channels)
        self.k_dropout = nn.Dropout(0.0)
        self.k_atten_norm = nn.LayerNorm(in_channels)
        # Post-Attention 
        self.post_atten = nn.MultiheadAttention(embed_dim=in_channels,
                                               num_heads=8,
                                               dropout=0.0)
        self.post_dropout = nn.Dropout(0.0)
        self.post_atten_norm = nn.LayerNorm(in_channels)
        # 
        self.ffn = FFN(in_channels,feedforfward_channels=2048,num_fcs=2)
        self.ffn_norm = nn.LayerNorm(in_channels)
        # 输出层
        self.pred = Pred()
        
        
    def forward(self,features,xxx):
        B,C,H,W = features.shape
        # features : b x c x h x w 
        # mask_preds: b x n x h x w 
        mask_preds = self.kernels(features)# 卷积操作
        # sigmoid处理并与阈值比较获得mask 
        # b x n x h x w 
        soft_sigmoid_mask = mask_preds.sigmoid()
        nonzero_inds = soft_sigmoid_mask > 0.5
        hard_sigmoid_mask = nonzero_inds.float()
        # V = r(A)r(S).T
        # b x n x c 
        V = torch.einsum('bnhw,bchw->bnc', hard_sigmoid_masks, features)
        # ------------------ 以上 Pre-Attentinon ------------------ # 
        
        # ----------------- 以下 DyconvAtten Block ---------------- #
        # 获取注意力的Q
        # n x c x k x k   k = 1
        proposal_kernels = self.kernels.weight.clone()
        # reshape  b x n x c x 1 x 1 -> b x n x (cxkxk)= b x n x c  与V的shape相同
        Q = proposal_kernels[None].expand(B,*proposal_kernels.size()).view().view(B,self.num_proposals,-1)
        # ------- 第一个 Block ------- #
        # b x n x c 
        f_point_out = self.f_atten(Q,V)
        V = V + self.f_dropout(f_point_out)
		# b x n x c 
        V = self.f_atten_norm(V)
        # ------- 第二个 Block ------- #
        k_point_out = self.k_atten(Q,V)
        V = V + self.k_dropout(k_point_out)
        # b x n x c
        V = self.k_atten_norm(V)
        
        # ------------------- 以上 DyconvAtten Block 结束 --------------- #
        
        # ------------------- Post Attention --------------- #
		K = V.permute(1,0,2)
        k_temp = self.post_atten(query=k,key=k,value=k)[0]
        k = k + self.post_dropout(k_temp)
        k = self.post_atten_norm(k.permute(1,0,2))
        # b x n x c -> b x n x c x k*k - > b x n x k*k x c = b x n x 1 x c
        obj_feat = k.view(B,self.num_proposals,self.in_channels,-1).permute(0,1,3,2)
        # b x n x k*k xc
        # 对应于图像中的Panoptic Kernels
        obj_feat = self.ffn_norm(self.ffn(obj_feat))
        # b x n x 1 x c
        cls_feat = obj_feat.sum(-2)
        mask_feat = obj_feat
        # b x n x k*k x c -> b x n x c x k*k ->  b x n x c x 1 x 1
        obj_feat = obj_feat.permute(0,1,3,2).view(B,self.num_propoasls,self.in_channels,self.conv_kernel_size_2d,self.conv_kernel_size_2d)
        # Pred
        cls_scores,new_mask_preds = self.pred(cls_feat,mask_feat,features)
        
        return cls_scores,new_mask_preds,obj_feat

        
class DySepConvAtten(nn.Module):
    def __init__(self,in_channels,kernel_size=3):
        super(DySepConvAtten,self).__init__()
        '''
        kernel_size : 一维卷积的卷积核大小
        '''
        self.kernel_size = kernel_size

        self.depth_weight_linear = nn.Linear(in_channels, self.kernel_size)
        self.point_weight_linear = nn.Linear(in_channels, self.num_proposals)
        self.norm = nn.LayerNorm(self.hidden_dim)
    def forward(self,Q,V):
        B = Q.shape[0]
        # b x n x 3
        dy_depth_conv_weight = self.depth_weight_linear(Q)
        # b x n x n
        dy_point_conv_weight = self.point_weight_linear(Q)
        # b x n x 1 x 3
        dy_depth_conv_weight = dy_depth_conv_weight.view(B,self.num_proposals,1,self.kernel_size)
        # b x n x n x 1
        dy_point_conv_weight = dy_point_conv_weight.view(B,self.num_proposals,self.num_proposals,1) 
        res = []
        # b x 1 x n x c
        V = V.unsqueeze(1)
        for i in range(B):# 依次对batch内每张图像进行处理
            # input: [1, N, C]
            # weight: [N, 1, K]
            # output: [1, N, C]
            out = F.relu(F.conv1d(input=value[i], weight=dy_depth_conv_weight[i], groups=N, padding="same"))
            # input: [1, N, C]
            # weight: [N, N, 1]
            # output: [1, N, C]
            out = F.conv1d(input=out, weight=dy_point_conv_weight[i], padding='same')
            res.append(out)
        # b x n x c
        point_out = torch.cat(res, dim=0)
        # b x n x c
        point_out = self.norm(point_out)
        return point_out

(3) Post-Attention模块

在后注意力模块中，采用多头自注意力(query=O,key=O,value=O)得到MO shape为[b,n,c], 然后 MO 和 前馈神经网络FFN 并通过来生成全景Kernels Output shape为[b,n,1,c],1=k*k,k=1.

注意该obj_feat, 在第二个block_T（后面会解释blockT的意义），将替代Proposal_Kernels的功能。

（4）预测输出

预测分两支： mask 预测 和 class 预测。

Post-Attention输出为Output shape[b,n,1,c]

<1> Class 预测

a. 先对Output 进行 FC + LayerNorm + ReLU操作，得到输出 Cls1 ,shape[b,n,c].

b. 再对Cls1 进行FC操作，得到输出 Cls,shape[b,n,cls_num+1]

<2> Mask 预测

a. 先对Output 进行 FC + LayerNorm + ReLU操作，得到输出 Mask1 ,shape[b,n,c].

b. 再对Mask1 进行FC操作，得到输出 Mask2,shape[b,n,c]

c. Mask2 联合 FPN的聚合特征S [b,c,h,w], 映射到Mask [b,n,h,w].

class Pred(nn.Module):
    def __init__(self,train_mode=True):
        super(Pred,self).__init__()
        self.num_cls_fcs = cfg.MODEL.YOSO.NUM_CLS_FCS
        self.num_mask_fcs = cfg.MODEL.YOSO.NUM_MASK_FCS
        self.num_classes = cfg.MODEL.YOSO.NUM_CLASSES
        
        self.train_mode = train_mode 
        self.cls_fcs = nn.ModuleList()
        
        for _ in range(self.num_cls_fcs):
            self.cls_fcs.append(nn.Linear(self.hidden_dim, self.hidden_dim, bias=False))
            self.cls_fcs.append(nn.LayerNorm(self.hidden_dim))
            self.cls_fcs.append(nn.ReLU(True))
        self.fc_cls = nn.Linear(self.hidden_dim, self.num_classes + 1)

        self.mask_fcs = nn.ModuleList()
        for _ in range(self.num_mask_fcs):
            self.mask_fcs.append(nn.Linear(self.hidden_dim, self.hidden_dim, bias=False))
            self.mask_fcs.append(nn.LayerNorm(self.hidden_dim))
            self.mask_fcs.append(nn.ReLU(True))
        self.fc_mask = nn.Linear(self.hidden_dim, self.hidden_dim)
   	def forward(self,cls_feat,mask_feat,features):
        '''
        cls_feat: b x n x 1 x c
        mask_feat:b x n x 1 x c 
        features: b x c x h x w
        '''
        if self.train_mode:
            # b x n x 1 x c 
            for cls_layer in self.cls_fcs:
                cls_feat = cls_layer(cls_feat)
            # b x n x 1 x (cls_num+1) -> b x n x cls_num+1
            cls_score = self.fc_cls(cls_feat).view(B, self.num_proposals, -1)
		else:
            cls_score = None 
		
        # b x n x 1 x c    1 = k * k
        for reg_layer in self.mask_fcs:
            mask_feat = reg_layer(mask_feat)
        # [B, N, K * K, C] -> [B, N, C]
        mask_kernels = self.fc_mask(mask_feat).squeeze(2)
        # features: b x c x h x w 
        # mask_kernels : b x n x c 
        # new_mask_preds : b x n x h x w  
        # 特征映射
        new_mask_preds = torch.einsum("bqc,bchw->bqhw", mask_kernels, features)
        return cls_score,new_mask_preds

4. 多个层堆叠

如上图所示，箭头的几个步骤为看作一个整体block_T，那么该block_T是可以堆叠多次的，以增加网络的深度，不过组要注意的是，从第二个block开始，其输入的 S 仍然为 FPN的输出features [b,c,h,w], 而 Q则为上面的obj_feat, mask_preds为上一个block的预测mask [b,n,c]。

更形象展开如下：