RT-DETR中的CCFF结构代码详解(Pytorch)

代码链接

lyuwenyu/RT-DETR: [CVPR 2024] Official RT-DETR (RTDETR paddle pytorch), Real-Time DEtection TRansformer, DETRs Beat YOLOs on Real-time Object Detection. 🔥 🔥 🔥 (github.com)https://github.com/lyuwenyu/RT-DETR

模型图

CCFF是作者提出的一种类似于特征金字塔的特征融合模块，S3，S4，S5是backbone的后三层，作者在论文中证明了只对S5进行尺度内交互，而不对更低级别的特征进行尺度内交互，并对次做法的合理性进行了证明，再次不多赘述

“基于上述分析，我们重新思考编码器的结构，提出了一种有效的混合编码器，由基于注意力的尺度内特征交互（AIFI）和基于 CNN 的跨尺度特征融合（CCFF）两个模块组成。具体来说，AIFI 通过使用单尺度 Transformer 编码器仅在 S5 上执行尺度内交互，进一步降低了基于变体 D 的计算成本。原因是将自注意力操作应用于具有更丰富语义概念的高级特征可以捕获概念实体之间的连接，这有助于后续模块对对象的定位和识别。然而，由于缺乏语义概念以及重复和与高级特征交互混淆的风险，低级特征的尺度内交互是不必要的。为了验证这一观点，我们仅在变体 D 中的 S5 上执行尺度内交互，实验结果如表 3 所示（参见第 DS5 行）。与D相比，DS5 不仅显着减少了延迟（快 35%），而且提高了准确性（AP 高 0.4%）。CCFF是基于跨尺度融合模块优化的，该模块将多个由卷积层组成的融合块插入到融合路径中。融合块的作用是将两个相邻的尺度特征融合到一个新的特征中，其结构如图5所示。融合块包含两个1 × 1卷积来调整通道数，还使用了N个RepBlock（由RepConv[8]组成的）进行特征融合，两条路径输出通过元素相加被融合。我们将混合编码器的计算表述为：”

本文要探讨的是这个CCFF结构，图像他的线画的有点乱，要是不看代码的话，完全理解不了，或者说，不敢吓理解，但是看过代码之后，你会发现他画的很正确：下面直入正题：

代码

首先直接找到Hybird Encoder类，CCFF结构被包含在内

@register
class HybridEncoder(nn.Module):
    def __init__(self,
                 in_channels=[512, 1024, 2048],#S3,S4,S5分别对应的通道数
                 feat_strides=[8, 16, 32],#如果对S3，S4，S5进行位置编码所需要的步长数字
                 hidden_dim=256,
                 nhead=8,
                 dim_feedforward = 1024,
                 dropout=0.0,
                 enc_act='gelu',
                 use_encoder_idx=[2],#一共传进来三层，编号是0，1，2，它只在最后一层进行AIFI（作者的Attention操作）所以这里就是个2
                 num_encoder_layers=1,
                 pe_temperature=10000,
                 expansion=1.0,
                 depth_mult=1.0,
                 act='silu',
                 eval_spatial_size=None):
        super().__init__()
        self.in_channels = in_channels
        self.feat_strides = feat_strides
        self.hidden_dim = hidden_dim
        self.use_encoder_idx = use_encoder_idx
        self.num_encoder_layers = num_encoder_layers
        self.pe_temperature = pe_temperature
        self.eval_spatial_size = eval_spatial_size

        self.out_channels = [hidden_dim for _ in range(len(in_channels))]
        self.out_strides = feat_strides

一个对三层特征的初始化映射层，以1*1卷积核将三层的通道数都映射为hidden_dim

        # channel projection
        self.input_proj = nn.ModuleList()
        for in_channel in in_channels:
            self.input_proj.append(
                nn.Sequential(
                    nn.Conv2d(in_channel, hidden_dim, kernel_size=1, bias=False),
                    nn.BatchNorm2d(hidden_dim)
                )
            )

下面是AIFI，其实就是一个单层的TransformerEncoderLayer，不多说

        # encoder transformer
        encoder_layer = TransformerEncoderLayer(
            hidden_dim, 
            nhead=nhead,
            dim_feedforward=dim_feedforward, 
            dropout=dropout,
            activation=enc_act)

        self.encoder = nn.ModuleList([
            TransformerEncoder(copy.deepcopy(encoder_layer), num_encoder_layers) for _ in range(len(use_encoder_idx))
        ])

一下是CCFF模块的定义，从上到下路径以及从下到上的路径

        # top-down fpn
        self.lateral_convs = nn.ModuleList()
        self.fpn_blocks = nn.ModuleList()
        for _ in range(len(in_channels) - 1, 0, -1):#2,1
            self.lateral_convs.append(ConvNormLayer(hidden_dim, hidden_dim, 1, 1, act=act))
            self.fpn_blocks.append(
                CSPRepLayer(hidden_dim * 2, hidden_dim, round(3 * depth_mult), act=act, expansion=expansion)
            )

        # bottom-up pan
        self.downsample_convs = nn.ModuleList()
        self.pan_blocks = nn.ModuleList()
        for _ in range(len(in_channels) - 1):
            self.downsample_convs.append(
                ConvNormLayer(hidden_dim, hidden_dim, 3, 2, act=act)
            )
            self.pan_blocks.append(
                CSPRepLayer(hidden_dim * 2, hidden_dim, round(3 * depth_mult), act=act, expansion=expansion)
            )

省略一部分位置编码的代码

快进到forward函数

forword函数首先验证输入特征的shape应该是[3,B,C,H,W],在将特征进行初始映射，再对S5进行尺度内交互，也就是将S5层过单层的TransformerEncoderLayer

代码如下：

    def forward(self, feats):
        assert len(feats) == len(self.in_channels)
        proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)]
        
        # encoder
        if self.num_encoder_layers > 0:#1
            for i, enc_ind in enumerate(self.use_encoder_idx):#实际上这里只有一次取值，就是i=0，enc_ind=2
                h, w = proj_feats[enc_ind].shape[2:]#获得backbone最后一层特征的高和宽分别存储在变量h和w里
                # flatten [B, C, H, W] to [B, HxW, C]
                src_flatten = proj_feats[enc_ind].flatten(2).permute(0, 2, 1)
                if self.training or self.eval_spatial_size is None:
                    pos_embed = self.build_2d_sincos_position_embedding(
                        w, h, self.hidden_dim, self.pe_temperature).to(src_flatten.device)
                else:
                    pos_embed = getattr(self, f'pos_embed{enc_ind}', None).to(src_flatten.device)

                memory = self.encoder[i](src_flatten, pos_embed=pos_embed)#经过Transformer的编码器，事实上是只有一层encoderlayer的encoder
                proj_feats[enc_ind] = memory.permute(0, 2, 1).reshape(-1, self.hidden_dim, h, w).contiguous()#经过编码之后再把它的形状恢复到之前的形状即[B,C,H,W],contiguous使连续存储
                # print([x.is_contiguous() for x in proj_feats ])

下面是最重点的CCFF流程，仔细看代码，注释很清楚，因此不再前解释了

# broadcasting and fusion
        #先从上到下，
        inner_outs = [proj_feats[-1]]#这个inner_outs存储的是最后一层经过Transformer编码器处理后的特征，与此同时从backbone提取的前两层都没有被处理过
        for idx in range(len(self.in_channels) - 1, 0, -1):#idx=2,1(不包括0）,初始为2
            feat_high = inner_outs[0]#在以上这种情况下，初始inner_outs中只有一个元素，所以又给最后一层经过TransformerEncoder过后的特征提出来了，即处理过后的proj_feats[-1]
            feat_low = proj_feats[idx - 1]#idx-1=1,0，初始为1，对应着S4层特征
            feat_high = self.lateral_convs[len(self.in_channels) - 1 - idx](feat_high)#len(self.in_channels)-1-idx=0,1初始为0,过一个1*1卷积(laterral_convs中全是1*1的带bn的卷积块)
            inner_outs[0] = feat_high#把feat_high存回去
            upsample_feat = F.interpolate(feat_high, scale_factor=2., mode='nearest')#feat_high[B,C,H,W]——>feat_high[B,C,2H,2W]
            inner_out = self.fpn_blocks[len(self.in_channels)-1-idx](torch.concat([upsample_feat, feat_low], dim=1))#len(self.in_channels)-1-idx:0,1,初始为0,初始将上采样的S5和S4特征在特征维度拼接起来送入CSP进行融合
            inner_outs.insert(0, inner_out)#将融合后的特征inner_out插入到列表的开头(索引 0 的位置)
            #该循环循环两遍
            #在第二次进行该循环的时候 feat_high是上一次循环融合好的特征 feat_low是S3的特征，将feat_high进行卷积上采样与feats_low进行融合，将结果再次插入到inner_outs中
            #经过两次循环之后 inner_outs中将包含三个元素 [0]是三层融合后的结构 [1]是S4和S5融合后的结果 [2]是S5自我卷积后的结果
        #再从下到上
        outs = [inner_outs[0]]#取的是三层从上到下融合后的结果再加一个维度(多套了一层[])
        for idx in range(len(self.in_channels) - 1):#len(self.in_channels)-1=2,idx=0,1
            feat_low = outs[-1]#取出三层从上到下融合后的特征
            feat_high = inner_outs[idx + 1]#idx+1=1,2;第一次循环取出S4和S5的融合结果
            downsample_feat = self.downsample_convs[idx](feat_low)#idx=0,1 采用卷积对低级特征进行下采样
            out = self.pan_blocks[idx](torch.concat([downsample_feat, feat_high], dim=1))#将下采样后的低级特征与高级特征在第通道维度上进行拼接，采用CSP进行融合
            outs.append(out)#将融合后的特征out拼接到outs的尾部
            #该循环会循环两边
            #在第二次进行该循环的时候 feat_low是S3和S4(经过从上到下处理)融合后结果，feats_high是S5的特征，将feat_low进行下采样与feat_high融合，将结果再次拼接到outs的末尾
            #经过两次循环时候 outs中包含三个元素 [0]是三个层自上而下融合后的结果 [1]是三个层自上而下融合后的结果再与上两层自上而下融合后的结果相容和的结果
            # [2]是五个层融合后的结果在和S5的初代卷积融合后的结果，称为六个层融合后的结果，也是S3,S4,S5自上而下再自下而上融合后的结果，也是S3,S4,S5各自被融合两次的结果

        return outs#最终返回outs数组