RT-DETR改进策略【Neck】| ArXiv 2023,基于U - Net v2中的的高效特征融合模块:SDI(Semantics and Detail Infusion)
一、本文介绍
本文聚焦于
利用U - Net v2中的SDI模块优化RT-DETR的目标检测网络模型
。
SDI
模块相较于传统模块独具特色,它融合了先进的
特征融合
思想,借助精心设计的结构,在确保计算资源高效利用的前提下,
巧妙地融合不同层级特征的语义信息与细节
,实现特征的全方位增强。在应用于
RT-DETR
的改进过程中,
SDI
模块能够助力模型
更精准地聚焦于图像中的目标物体,有效抑制背景及其他干扰因素,凸显目标的关键特征与位置信息
,从而大幅提升
RT-DETR
在复杂场景下的目标检测精度与稳定性。
二、SDI介绍
U-NET V2: RETHINKING THE SKIP CONNECTIONS OF U-NET FOR MEDICAL IMAGE SEGMENTATION
U-Net V2
中的
SDI模块
在整个网络架构中起着关键作用,其设计旨在
解决传统模型在特征融合方面的不足
,通过独特的结构原理实现更高效的语义信息和细节融合,从而提升医学图像分割的性能。
2.1 设计出发点
在传统的
U - Net模型
中,基于
U - Net
的模型连接方式
在整合低层次和高层次特征时可能效果不佳
。
- 编码器提取的低层次特征通常保留更多细节但缺乏足够语义信息且可能含噪声,高层次特征虽有较多语义信息却因分辨率降低缺少精确细节。
- 简单的拼接融合依赖网络学习能力,在医学影像数据有限的情况下是个挑战,且会增加GPU内存消耗和计算量。
因此,需要一种更有效的方式来融合不同层次的特征,这就是
SDI模块
的设计出发点。
2.2 SDI结构原理
- 首先,对编码器生成的每一层级(i)的特征 f i 0 f_{i}^{0} f i 0 应用空间和通道注意力机制 φ i s \varphi_{i}^{s} φ i s 和 ϕ i c \phi_{i}^{c} ϕ i c ,公式为 f i 1 = ϕ i c ( φ i s ( f i 0 ) ) f_{i}^{1}=\phi_{i}^{c}\left(\varphi_{i}^{s}\left(f_{i}^{0}\right)\right) f i 1 = ϕ i c ( φ i s ( f i 0 ) ) 使特征 能够整合局部空间信息和全局通道信息 。然后应用 1 × 1 1\times1 1 × 1 卷积将 f i 1 f_{i}^{1} f i 1 的通道数减少到 c c c (超参数),得到 f i 2 f_{i}^{2} f i 2 。
- 在解码器的每一层级 i i i ,以 f i 2 f_{i}^{2} f i 2 为目标参考,调整其他层级特征图的大小使其与 f i 2 f_{i}^{2} f i 2 分辨率匹配,对于 j < i j < i j < i : f i j 3 = D ( f j 2 , ( H i , W i ) ) f_{ij}^{3}=D\left(f_{j}^{2},\left(H_{i}, W_{i}\right)\right) f ij 3 = D ( f j 2 , ( H i , W i ) ) ( D D D 为自适应平均池化); j = i j = i j = i 时: f i j 3 = ( I ) ( f j 2 ) f_{ij}^{3}=(I)\left(f_{j}^{2}\right) f ij 3 = ( I ) ( f j 2 ) ( I I I 为恒等映射); j > i j > i j > i 时: f i j 3 = U ) ( f j 2 , ( H i , W i ) ) f_{ij}^{3}=U)\left(f_{j}^{2},\left(H_{i}, W_{i}\right)\right) f ij 3 = U ) ( f j 2 , ( H i , W i ) ) ( U U U 为双线性插值)。接着对调整后的特征图 f i j 3 f_{ij}^{3} f ij 3 应用 3 × 3 3\times3 3 × 3 卷积进行平滑,得到 f i j 4 = θ i j ( f i j 3 ) f_{ij}^{4}=\theta_{ij}\left(f_{ij}^{3}\right) f ij 4 = θ ij ( f ij 3 ) 。最后,将所有调整为相同分辨率的第(i)层级特征图通过逐元素哈达玛积(H)进行融合,公式为 f i 5 = H ( [ f i 1 4 , f i 2 4 , ⋯ , f i M 4 ] ) f_{i}^{5}=H\left(\left[f_{i1}^{4}, f_{i2}^{4},\cdots, f_{iM}^{4}\right]\right) f i 5 = H ( [ f i 1 4 , f i 2 4 , ⋯ , f i M 4 ] ) 并将 f i 5 f_{i}^{5} f i 5 发送到第 i i i 层级解码器进行后续处理。
2.3 优势
- 从实验结果来看,在皮肤病变分割和息肉分割数据集上,U - Net V2相比其他先进方法有更好的分割效果。
- 在计算复杂度、GPU内存使用和推理时间方面,U - Net V2也表现出优势。相比UNet++,U - Net V2引入的参数更少,GPU内存使用量更低,FLOPs和FPS更优,说明 SDI模块在提升性能的同时没有带来过多的计算和存储负担。
论文: https://arxiv.org/pdf/2311.17791
源码: https://github.com/yaoppeng/U-Net_v2
三、SDI的实现代码
SDI模块
的实现代码如下:
import torch
from torch import nn
import torch.nn.functional as F
def autopad(k, p=None, d=1):
"""
Pads kernel to 'same' output shape, adjusting for optional dilation; returns padding size.
`k`: kernel, `p`: padding, `d`: dilation.
"""
if d > 1:
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
class Conv(nn.Module):
# Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
default_act = nn.SiLU() # default activation
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
"""Initializes a standard convolution layer with optional batch normalization and activation."""
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self, x):
"""Applies a convolution followed by batch normalization and an activation function to the input tensor `x`."""
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
"""Applies a fused convolution and activation function to the input tensor `x`."""
return self.act(self.conv(x))
class GSConv(nn.Module):
# GSConv https://github.com/AlanLi1997/slim-neck-by-gsconv
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
super().__init__()
c_ = c2 // 2
self.cv1 = Conv(c1, c_, k, s, p, g, d, Conv.default_act)
self.cv2 = Conv(c_, c_, 5, 1, p, c_, d, Conv.default_act)
def forward(self, x):
x1 = self.cv1(x)
x2 = torch.cat((x1, self.cv2(x1)), 1)
# shuffle
# y = x2.reshape(x2.shape[0], 2, x2.shape[1] // 2, x2.shape[2], x2.shape[3])
# y = y.permute(0, 2, 1, 3, 4)
# return y.reshape(y.shape[0], -1, y.shape[3], y.shape[4])
b, n, h, w = x2.size()
b_n = b * n // 2
y = x2.reshape(b_n, 2, h * w)
y = y.permute(1, 0, 2)
y = y.reshape(2, -1, n // 2, h, w)
return torch.cat((y[0], y[1]), 1)
class SDI(nn.Module):
def __init__(self, channels):
super().__init__()
# self.convs = nn.ModuleList([nn.Conv2d(channel, channels[0], kernel_size=3, stride=1, padding=1) for channel in channels])
self.convs = nn.ModuleList([GSConv(channel, channels[0]) for channel in channels])
def forward(self, xs):
ans = torch.ones_like(xs[0])
target_size = xs[0].shape[2:]
for i, x in enumerate(xs):
if x.shape[-1] > target_size[-1]:
x = F.adaptive_avg_pool2d(x, (target_size[0], target_size[1]))
elif x.shape[-1] < target_size[-1]:
x = F.interpolate(x, size=(target_size[0], target_size[1]),
mode='bilinear', align_corners=True)
ans = ans * self.convs[i](x)
return ans
四、添加步骤
4.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
SDI.py
,将
第三节
中的代码粘贴到此处
4.2 修改二
在
AddModules
文件夹下新建
__init__.py
(已有则不用新建),在文件内导入模块:
from .SDI import *
4.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中,需要在两处位置添加各模块类名称。
首先:导入模块
然后,在
parse_model函数
中添加如下代码:
elif m is SDI:
args = [[ch[x] for x in f]]
五、yaml模型文件
5.1 模型改进版本⭐
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例,在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-SDI.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-SDI.yaml
文件下,修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
颈部网络
中的
Concat
替换成
SDI模块
。
# 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-P4/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-P5/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, SDI, []]
- [-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, SDI, []] # 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, SDI, []] # cat Y4
- [-1, 3, RepC3, [256]] # F4 (24), pan_blocks.0
- [-1, 1, Conv, [256, 3, 2]] # 25, downsample_convs.1
- [[-1, 12], 1, SDI, []] # cat Y5
- [-1, 3, RepC3, [256]] # F5 (27), pan_blocks.1
- [[21, 24, 27], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5)
六、成功运行结果
打印网络模型可以看到
SDI
已经加入到模型中,并可以进行训练了。
rtdetr-l-SDI :
rtdetr-l-SDI summary: 742 layers, 32,575,683 parameters, 32,575,683 gradients, 106.9 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 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
14 7 1 262656 ultralytics.nn.modules.conv.Conv [1024, 256, 1, 1, None, 1, 1, False]
15 [-2, -1] 1 72960 ultralytics.nn.AddModules.SDI.SDI [[256, 256]]
16 -1 3 2101248 ultralytics.nn.modules.block.RepC3 [256, 256, 3]
17 -1 1 66048 ultralytics.nn.modules.conv.Conv [256, 256, 1, 1]
18 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
19 3 1 131584 ultralytics.nn.modules.conv.Conv [512, 256, 1, 1, None, 1, 1, False]
20 [-2, -1] 1 72960 ultralytics.nn.AddModules.SDI.SDI [[256, 256]]
21 -1 3 2101248 ultralytics.nn.modules.block.RepC3 [256, 256, 3]
22 -1 1 590336 ultralytics.nn.modules.conv.Conv [256, 256, 3, 2]
23 [-1, 17] 1 72960 ultralytics.nn.AddModules.SDI.SDI [[256, 256]]
24 -1 3 2101248 ultralytics.nn.modules.block.RepC3 [256, 256, 3]
25 -1 1 590336 ultralytics.nn.modules.conv.Conv [256, 256, 3, 2]
26 [-1, 12] 1 72960 ultralytics.nn.AddModules.SDI.SDI [[256, 256]]
27 -1 3 2101248 ultralytics.nn.modules.block.RepC3 [256, 256, 3]
28 [21, 24, 27] 1 7303907 ultralytics.nn.modules.head.RTDETRDecoder [1, [256, 256, 256]]
rtdetr-l-SDI summary: 742 layers, 32,575,683 parameters, 32,575,683 gradients, 106.9 GFLOPs