RT-DETR改进策略【Neck】| 使用CARAFE轻量级通用上采样算子

一、本文介绍

本文记录的是利用CARAFE上采样对RT-DETR的颈部网络进行改进的方法研究。 RT-DETR 采用传统的最近邻插值的方法，仅考虑子像素邻域，无法捕获密集预测任务所需的丰富语义信息，从而影响模型在密集预测任务中的性能。 CARAFE 通过在大感受野内聚合信息、能够实时适应实例特定内容且保持计算效率。

二、CARAFE介绍

CARAFE ： 内容感知的特征重新组合

CARAFE（Content - Aware ReAssembly of FEatures） 是一种用于特征上采样的操作符，其设计原理和优势如下：

2.1、CARAFE原理

总体框架 ： CARAFE 由两个关键组件组成，即 核预测模块 和 内容感知重组模块 。给定大小为 $C \times H \times W$ 的特征图 $x$ 和上采样比例 $\sigma$ （假设 $\sigma$ 是整数）， CARAFE 将产生大小为 $C×\sigma H×\sigma W$ 的新特征图 $X^{'}$ 。
核预测模块 ：负责以内容感知的方式生成重组核。每个源位置 $x$ 上对应 $\sigma^{2}$ 个目标位置 $X^{'}$ 上，每个目标位置需要一个 $k_{u p}×k_{u p}$ 的重组核，因此该模块将输出大小为 $C_{u p}×H×W$ 的重组核，其中 $C_{u p} = \sigma^{2}k_{u p}^{2}$ 。该模块由三个子模块组成，分别是通道压缩器、内容编码器和核归一化器。
- 通道压缩器 ：采用 $1 \times 1$ 卷积层将输入特征通道从 $C$ 压缩到 $C_{m}$ ，减少特征图的通道数，从而减少后续步骤的参数和计算成本，使CARAFE更高效。
- 内容编码器 ：使用核大小为 $k_{encoder}$ 的卷积层根据输入特征的内容生成重组核，编码器的参数为 $k_{encoder}×k_{encoder}×C_{m}×C_{u p}$ 。增大 $k_{encoder}$ 可以扩大编码器的感受野，利用更大区域内的上下文信息，但计算复杂度也会随之增加。通过研究， $k_{encoder} = k_{u p} - 2$ 是性能和效率之间的良好平衡。
- 核归一化器 ：在将每个 $k_{u p}×k_{u p}$ 重组核应用于输入特征图之前，使用softmax函数对其进行空间归一化，使核值之和为 $1$ ，这是对局部区域的软选择。
内容感知重组模块 ：对于目标位置 $l^{'}$ 和以 $l = (i, j)$ 为中心的相应方形区域 $N(X_{l}, k_{u p})$ ，重组过程如公式 $\mathcal{X}_{l'}' = \sum_{n = -r}^{r}\sum_{m = -r}^{r}\mathcal{W}_{l'(n, m)}·\mathcal{X}_{(i + n, j + m)}$ 所示，其中 $\lfloor k_{u p} / 2\rfloor$ 。通过重组核，区域 $N(X_{l}, k_{u p})$ 中的每个像素根据特征的内容而不是位置的距离对上采样像素 $l^{'}$ 的贡献不同，使得重组后的特征图的语义比原始特征图更强。

在这里插入图片描述

2.2、优势

大视野 ：与之前仅利用子像素邻域的工作不同， CARAFE 可以在大感受野内聚合上下文信息。
内容感知处理 ： CARAFE 能够根据实例特定的内容进行实时处理，为每个样本生成自适应的内核，而不是使用固定的内核。
轻量且计算快速 ： CARAFE 引入的计算开销很小，可以很容易地集成到现代网络架构中。
通用性和有效性 ：在对象检测、实例分割、语义分割、图像修复等广泛的密集预测任务中，CARAFE都能显著提升性能。

论文： https://arxiv.org/abs/1905.02188
源码： https://github.com/tiny-smart/dysample

三、CARAFE的实现代码

CARAFE模块 的实现代码如下：

class CARAFE(nn.Module):
    def __init__(self, c, k_enc=3, k_up=5, c_mid=64, scale=2):
        """ The unofficial implementation of the CARAFE module.
        The details are in "https://arxiv.org/abs/1905.02188".
        Args:
            c: The channel number of the input and the output.
            c_mid: The channel number after compression.
            scale: The expected upsample scale.
            k_up: The size of the reassembly kernel.
            k_enc: The kernel size of the encoder.
        Returns:
            X: The upsampled feature map.
        """
        super(CARAFE, self).__init__()
        self.scale = scale
 
        self.comp = Conv(c, c_mid)
        self.enc = Conv(c_mid, (scale * k_up) ** 2, k=k_enc, act=False)
        self.pix_shf = nn.PixelShuffle(scale)
 
        self.upsmp = nn.Upsample(scale_factor=scale, mode='nearest')
        self.unfold = nn.Unfold(kernel_size=k_up, dilation=scale,
                                padding=k_up // 2 * scale)
 
    def forward(self, X):
        b, c, h, w = X.size()
        h_, w_ = h * self.scale, w * self.scale
 
        W = self.comp(X)  # b * m * h * w
        W = self.enc(W)  # b * 100 * h * w
        W = self.pix_shf(W)  # b * 25 * h_ * w_
        W = torch.softmax(W, dim=1)  # b * 25 * h_ * w_
 
        X = self.upsmp(X)  # b * c * h_ * w_
        X = self.unfold(X)  # b * 25c * h_ * w_
        X = X.view(b, c, -1, h_, w_)  # b * 25 * c * h_ * w_
 
        X = torch.einsum('bkhw,bckhw->bchw', [W, X])  # b * c * h_ * w_
        return X

四、添加步骤

4.1 修改ultralytics/nn/modules/block.py

此处需要修改的文件是 ultralytics/nn/modules/block.py

block.py中定义了网络结构的通用模块 ，我们想要加入新的模块就只需要将模块代码放到这个文件内即可。

将 CARAFE 模块代码添加到此文件下。

在这里插入图片描述

4.2 修改ultralytics/nn/modules/ init .py

此处需要修改的文件是 ultralytics/nn/modules/__init__.py

__init__.py 文件中定义了所有模块的初始化，我们只需要将 block.py 中的新的模块命添加到对应的函数即可。

CARAFE 在 block.py 中实现，所有要添加在 from .block import ：

from .block import (
    C1,
    C2,
    ...
    CARAFE
)

在这里插入图片描述

4.3 修改ultralytics/nn/modules/tasks.py

在 tasks.py 文件中，需要在两处位置添加各模块类名称。

首先：在函数声明中引入 CARAFE

在这里插入图片描述

其次：在 parse_model函数 中注册 CARAFE 模块

在这里插入图片描述

elif m in [CARAFE]:
    args = [ch[f], *args[0:]]

在这里插入图片描述

五、yaml模型文件

5.1 模型改进⭐

在代码配置完成后，配置模型的YAML文件。

此处以 ultralytics/cfg/models/rt-detr/rtdetr-l.yaml 为例，在同目录下创建一个用于自己数据集训练的模型文件 rtdetr-l-CARAFE.yaml 。

将 rtdetr-l.yaml 中的内容复制到 rtdetr-l-CARAFE.yaml 文件下，修改 nc 数量等于自己数据中目标的数量。

📌 修改方法是将 CARAFE模块 替换 rtdetr-l 颈部网络中的上采样模块 nn.Upsample 。

# Ultralytics YOLO 🚀, AGPL-3.0 license
# RT-DETR-l object detection model with P3-P5 outputs. For details see https://docs.ultralytics.com/models/rtdetr

# Parameters
nc: 1 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n-cls.yaml' will call yolov8-cls.yaml with scale 'n'
  # [depth, width, max_channels]
  l: [1.00, 1.00, 1024]

backbone:
  # [from, repeats, module, args]
  - [-1, 1, HGStem, [32, 48]] # 0-P2/4
  - [-1, 6, HGBlock, [48, 128, 3]] # stage 1

  - [-1, 1, DWConv, [128, 3, 2, 1, False]] # 2-P3/8
  - [-1, 6, HGBlock, [96, 512, 3]] # stage 2

  - [-1, 1, DWConv, [512, 3, 2, 1, False]] # 4-P3/16
  - [-1, 6, HGBlock, [192, 1024, 5, True, False]] # cm, c2, k, light, shortcut
  - [-1, 6, HGBlock, [192, 1024, 5, True, True]]
  - [-1, 6, HGBlock, [192, 1024, 5, True, True]] # stage 3

  - [-1, 1, DWConv, [1024, 3, 2, 1, False]] # 8-P4/32
  - [-1, 6, HGBlock, [384, 2048, 5, True, False]] # stage 4

head:
  - [-1, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 10 input_proj.2
  - [-1, 1, AIFI, [1024, 8]]
  - [-1, 1, Conv, [256, 1, 1]] # 12, Y5, lateral_convs.0

  - [-1, 1, CARAFE, []]
  - [7, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 14 input_proj.1
  - [[-2, -1], 1, Concat, [1]]
  - [-1, 3, RepC3, [256]] # 16, fpn_blocks.0
  - [-1, 1, Conv, [256, 1, 1]] # 17, Y4, lateral_convs.1

  - [-1, 1, CARAFE, []]
  - [3, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 19 input_proj.0
  - [[-2, -1], 1, Concat, [1]] # cat backbone P4
  - [-1, 3, RepC3, [256]] # X3 (21), fpn_blocks.1

  - [-1, 1, Conv, [256, 3, 2]] # 22, downsample_convs.0
  - [[-1, 17], 1, Concat, [1]] # cat Y4
  - [-1, 3, RepC3, [256]] # F4 (24), pan_blocks.0

  - [-1, 1, Conv, [256, 3, 2]] # 25, downsample_convs.1
  - [[-1, 12], 1, Concat, [1]] # cat Y5
  - [-1, 3, RepC3, [256]] # F5 (27), pan_blocks.1

  - [[21, 24, 27], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5)

六、成功运行结果

分别打印网络模型可以看到 CARAFE模块 已经加入到模型中，并可以进行训练了。

rtdetr-l-CARAFE ：

rtdetr-l-CARAFE summary: 702 layers, 32,956,755 parameters, 32,956,755 gradients, 108.3 GFLOPs

                   from  n    params  module                                       arguments                     
  0                  -1  1     25248  ultralytics.nn.modules.block.HGStem          [3, 32, 48]                   
  1                  -1  6    155072  ultralytics.nn.modules.block.HGBlock         [48, 48, 128, 3, 6]           
  2                  -1  1      1408  ultralytics.nn.modules.conv.DWConv           [128, 128, 3, 2, 1, False]    
  3                  -1  6    839296  ultralytics.nn.modules.block.HGBlock         [128, 96, 512, 3, 6]          
  4                  -1  1      5632  ultralytics.nn.modules.conv.DWConv           [512, 512, 3, 2, 1, False]    
  5                  -1  6   1695360  ultralytics.nn.modules.block.HGBlock         [512, 192, 1024, 5, 6, True, False]
  6                  -1  6   2055808  ultralytics.nn.modules.block.HGBlock         [1024, 192, 1024, 5, 6, True, True]
  7                  -1  6   2055808  ultralytics.nn.modules.block.HGBlock         [1024, 192, 1024, 5, 6, True, True]
  8                  -1  1     11264  ultralytics.nn.modules.conv.DWConv           [1024, 1024, 3, 2, 1, False]  
  9                  -1  6   6708480  ultralytics.nn.modules.block.HGBlock         [1024, 384, 2048, 5, 6, True, False]
 10                  -1  1    524800  ultralytics.nn.modules.conv.Conv             [2048, 256, 1, 1, None, 1, 1, False]
 11                  -1  1    789760  ultralytics.nn.modules.transformer.AIFI      [256, 1024, 8]                
 12                  -1  1     66048  ultralytics.nn.modules.conv.Conv             [256, 256, 1, 1]              
 13                  -1  1     74312  ultralytics.nn.AddModules.CARAFE.CARAFE      [256]                         
 14                   7  1    262656  ultralytics.nn.modules.conv.Conv             [1024, 256, 1, 1, None, 1, 1, False]
 15            [-2, -1]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 16                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 17                  -1  1     66048  ultralytics.nn.modules.conv.Conv             [256, 256, 1, 1]              
 18                  -1  1     74312  ultralytics.nn.AddModules.CARAFE.CARAFE      [256]                         
 19                   3  1    131584  ultralytics.nn.modules.conv.Conv             [512, 256, 1, 1, None, 1, 1, False]
 20            [-2, -1]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 21                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 22                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
 23            [-1, 17]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 24                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 25                  -1  1    590336  ultralytics.nn.modules.conv.Conv             [256, 256, 3, 2]              
 26            [-1, 12]  1         0  ultralytics.nn.modules.conv.Concat           [1]                           
 27                  -1  3   2232320  ultralytics.nn.modules.block.RepC3           [512, 256, 3]                 
 28        [21, 24, 27]  1   7303907  ultralytics.nn.modules.head.RTDETRDecoder    [1, [256, 256, 256]]          
rtdetr-l-CARAFE summary: 702 layers, 32,956,755 parameters, 32,956,755 gradients, 108.3 GFLOPs

学习资源站

RT-DETR改进策略【Neck】使用CARAFE轻量级通用上采样算子_carafe上采样代码-

RT-DETR改进策略【Neck】| 使用CARAFE轻量级通用上采样算子

一、本文介绍

二、CARAFE介绍

2.1、CARAFE原理

2.2、优势

三、CARAFE的实现代码

四、添加步骤

4.1 修改ultralytics/nn/modules/block.py

4.2 修改ultralytics/nn/modules/ init .py

4.3 修改ultralytics/nn/modules/tasks.py

五、yaml模型文件

5.1 模型改进⭐

六、成功运行结果