RT-DETR融合[ECCV2024]自调制特征聚合SMFA模块及相关改进思路

news2024/12/29 4:29:12

RT-DETR使用教程： RT-DETR使用教程

RT-DETR改进汇总贴：RT-DETR更新汇总贴

《SMFANet: A Lightweight Self-Modulation Feature Aggregation Network for Efficient Image Super-Resolution》

一、模块介绍

论文链接：https://link.springer.com/chapter/10.1007/978-3-031-72973-7_21

代码链接：https://github.com/Zheng-MJ/SMFANet?tab=readme-ov-file

论文速览：

基于 Transformer 的修复方法取得了显着的性能，因为 Transformer 的自注意力（SA）可以探索非局部信息以获得更好的高分辨率图像重建。然而，关键的点积 SA 需要大量的计算资源。此外，SA 机制的低通特性限制了其捕获局部细节的能力，从而导致平滑的重建结果。为了解决这些问题，作者提出了一个自调制特征聚合（SMFA）模块，以协同利用局部和非局部特征交互来实现更准确的重建。具体来说，SMFA 模块采用高效的自我注意近似（EASA）分支来对非局部信息进行建模，并使用局部细节估计（LDE）分支来捕获局部细节。此外，作者进一步引入了基于部分卷积的前馈网络（PCFN）来改进从 SMFA 派生的代表性特征。大量实验表明，所提出的 SMFANet 系列在公共基准数据集上实现了更好的重建性能和计算效率之间的权衡。特别是，与×4 SwinIR-light，SMFANet+ 在五个公共测试集中平均实现了 0.14 dB 的性能提升，并且×运行速度提高 10 倍，模型复杂度仅为 43% 左右（例如 FLOPs）。

总结：一种基于自调制特征聚合模块（SMFA）的高分辨率图像重建方法，实测与其他模块融合有提升。

二、加入到RT-DETR中

2.1 创建脚本文件

首先在ultralytics->nn路径下创建blocks.py脚本，用于存放模块代码。

2.2 复制代码

复制代码粘到刚刚创建的blocks.py脚本中，如下图所示：

import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
class DMlp(nn.Module):
    def __init__(self, dim, growth_rate=2.0):
        super().__init__()
        hidden_dim = int(dim * growth_rate)
        self.conv_0 = nn.Sequential(
            nn.Conv2d(dim,hidden_dim,3,1,1,groups=dim),
            nn.Conv2d(hidden_dim,hidden_dim,1,1,0)
        )
        self.act =nn.GELU()
        self.conv_1 = nn.Conv2d(hidden_dim, dim, 1, 1, 0)
 
    def forward(self, x):
        x = self.conv_0(x)
        x = self.act(x)
        x = self.conv_1(x)
        return x
 
 
class SMFA(nn.Module):
    def __init__(self, dim=36):
        super(SMFA, self).__init__()
        self.linear_0 = nn.Conv2d(dim,dim*2,1,1,0)
        self.linear_1 = nn.Conv2d(dim,dim,1,1,0)
        self.linear_2 = nn.Conv2d(dim,dim,1,1,0)
 
        self.lde = DMlp(dim,2)
 
        self.dw_conv = nn.Conv2d(dim,dim,3,1,1,groups=dim)
 
        self.gelu = nn.GELU()
        self.down_scale = 8
 
        self.alpha = nn.Parameter(torch.ones((1,dim,1,1)))
        self.belt = nn.Parameter(torch.zeros((1,dim,1,1)))
 
    def forward(self, f):
        _,_,h,w = f.shape
        y, x = self.linear_0(f).chunk(2, dim=1)
        x_s = self.dw_conv(F.adaptive_max_pool2d(x, (h // self.down_scale, w // self.down_scale)))
        x_v = torch.var(x, dim=(-2,-1), keepdim=True)
        x_l = x * F.interpolate(self.gelu(self.linear_1(x_s * self.alpha + x_v * self.belt)), size=(h,w), mode='nearest')
        y_d = self.lde(y)
        return self.linear_2(x_l + y_d)

2.3 更改task.py文件

打开ultralytics->nn->modules->task.py，在脚本空白处导入函数。

from ultralytics.nn.blocks import *

之后找到模型解析函数parse_model（约在tasks.py脚本中940行左右位置，可能因代码版本不同变动），在该函数的最后一个else分支上面增加相关解析代码。

        elif m is SMFA:
            c2 = ch[f]
            args = [ch[f]]

2.4 更改yaml文件

yam文件解读：YOLO系列 “.yaml“文件解读_yolo yaml文件-CSDN博客

打开更改ultralytics/cfg/models/rt-detr路径下的rtdetr-l.yaml文件，替换原有模块。（放在该位置仅能插入该模块，具体效果未知。博主精力有限，仅完成与其他模块二次创新融合的测试，结构图见文末，代码见群文件更新。）

# 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: 80 # 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, 2, SMFA, []] # 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, nn.Upsample, [None, 2, "nearest"]]
  - [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, nn.Upsample, [None, 2, "nearest"]]
  - [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)

2.5 修改train.py文件

创建Train_RT脚本用于训练。

from ultralytics.models import RTDETR
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

if __name__ == '__main__':
    model = RTDETR(model='ultralytics/cfg/models/rt-detr/rtdetr-l.yaml')
    # model.load('yolov8n.pt')
    model.train(data='./data.yaml', epochs=2, batch=1, device='0', imgsz=640, workers=2, cache=False,
                amp=True, mosaic=False, project='runs/train', name='exp')