RT-DETR改进策略【注意力机制篇】| ICLR2023 高效计算与全局局部信息融合的 Sea_Attention 模块（含HGBlock二次创新）
一、本文介绍
本文记录的是
利用
SeaFormer++
模型中提出的
Sea_Attention
模块优化
RT-DETR
的目标检测网络模型
。
Sea_Attention
利用
挤压轴向注意力
有效地提取
全局语义信息
，并通过
细节增强核
补充局部细节
，优化了
Transformer块
的特征提取能力。本文将其加入到
RT-DETR
的不同位置中，使模型能够在不引入过多计算开销的情况下聚合空间信息。
二、混合局部通道注意力介绍
SeaFormer++: Squeeze-enhanced Axial Transformer for Mobile Visual Recognition
2.1 出发点
传统的全局自注意力机制在处理高分辨率图像时计算成本和内存需求高，不适合移动设备。为了解决这个问题，需要设计一种高效的注意力模块，既能提取全局语义信息，又能保持较低的计算复杂度和内存占用，同时还要能补充局部细节信息，以满足移动设备上的语义分割任务需求。
2.2 原理
2.2.1 Squeeze Axial attention（挤压轴向注意力）
- 通过自适应地将输入特征图在水平和垂直方向上进行挤压操作，将每个轴上的所有标记映射到一个单一的标记，从而以一种自适应的方式将全局信息保留在单个轴上。
- 在水平方向上，使用可学习的掩码将查询的所有标记映射到每行的一个单一标记；在垂直方向上同理。这样可以大大降低计算复杂度。
- 同时，为了使挤压后的轴向注意力具有位置感知能力，引入位置嵌入，使得挤压后的查询和键能够感知到它们在挤压后的轴向特征中的位置。
2.2.2 Detail enhancement kernel（细节增强核）
- 由于挤压操作会牺牲局部细节，因此使用一个基于卷积的辅助核来增强空间细节。
- 首先从输入特征图中获取查询、键和值，并在通道维度上进行拼接，然后通过一个由3×3深度可分离卷积和批量归一化组成的块来聚合辅助局部细节。
- 最后通过一个线性投影和激活函数以及批量归一化将维度压缩到指定维度，并生成细节增强权重，与挤压轴向注意力得到的特征进行融合。
2.3 结构
2.3.1 Squeeze Axial attention部分
- 包括水平和 垂直方向 的挤压操作，以及相应的 位置嵌入 。
- 对于水平方向，计算 q ( h ) = 1 W ( q → ( H , C q k , W ) A W → ( H , W , 1 ) ) → ( H , C q k ) q_{(h)}=\frac{1}{W}\left(q^{\to\left(H, C_{qk}, W\right)} A_{W}^{\to(H, W, 1)}\right)^{\to\left(H, C_{qk}\right)} q ( h ) = W 1 ( q → ( H , C q k , W ) A W → ( H , W , 1 ) ) → ( H , C q k ) ，其中 q q q 是从输入特征图 x x x 通过线性投影得到的查询， A W A_{W} A W 是可学习的掩码，通过在输入特征图上应用1×1卷积和批量归一化层得到。垂直方向同理计算 q ( v ) q_{(v)} q ( v ) 。
- 位置嵌入通过从可学习参数线性插值得到，分别为 r ( h ) q , r ( h ) k ∈ R H × C q k r_{(h)}^{q}, r_{(h)}^{k} \in \mathbb{R}^{H ×C_{qk}} r ( h ) q , r ( h ) k ∈ R H × C q k 和 r ( v ) q , r ( v ) k ∈ R W × C q k r_{(v)}^{q}, r_{(v)}^{k} \in \mathbb{R}^{W ×C_{qk}} r ( v ) q , r ( v ) k ∈ R W × C q k ，并应用到相应的挤压后的查询和键上。
2.3.2 Detail enhancement kernel部分
- 从输入特征图(x)通过另一组线性投影获取查询、键和值 W q ( e ) , W k ( e ) ∈ R C q k × C , W v ( e ) ∈ R C v × C W_{q}^{(e)}, W_{k}^{(e)} \in \mathbb{R}^{C_{qk} ×C}, W_{v}^{(e)} \in \mathbb{R}^{C_{v} ×C} W q ( e ) , W k ( e ) ∈ R C q k × C , W v ( e ) ∈ R C v × C ，然后在通道维度上拼接并通过3×3深度可分离卷积和批量归一化块，再经过线性投影、激活函数和批量归一化得到细节增强权重，与挤压轴向注意力的结果融合。
2.4 优势
- 计算高效 ：通过挤压轴向注意力将计算复杂度从 O ( ( H + W ) H W ) O((H + W)HW) O (( H + W ) H W ) 降低到 O ( H W ) O(HW) O ( H W ) ，同时通过自适应的挤压和扩展操作，在不引入过多计算开销的情况下聚合空间信息。
- 有效提取全局和局部信息 ：既能通过挤压轴向注意力有效地提取全局语义信息，又能通过细节增强核补充局部细节，优化了Transformer块的特征提取能力。
- 适合移动设备 ：只包含卷积、池化、矩阵乘法等移动设备友好的操作，在移动设备上能够实现高效的语义分割。
论文： https://arxiv.org/pdf/2301.13156
源码： https://github.com/fudan-zvg/SeaFormer
三、Sea_Attention的实现代码
Sea_Attention
及其改进的实现代码如下：
import math
import torch
from torch import nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule
from mmcv.cnn import build_norm_layer
from timm.models.registry import register_model
from ultralytics.nn.modules.conv import LightConv
def autopad(k, p=None, d=1): # kernel, padding, dilation
# Pad to 'same' shape outputs
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):
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):
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
return self.act(self.conv(x))
def _make_divisible(v, divisor, min_value=None):
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by 8
It can be seen here:
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
:param v:
:param divisor:
:param min_value:
:return:
"""
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_v < 0.9 * v:
new_v += divisor
return new_v
def drop_path(x, drop_prob: float = 0., training: bool = False):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
'survival rate' as the argument.
"""
if drop_prob == 0. or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_() # binarize
output = x.div(keep_prob) * random_tensor
return output
class DropPath(nn.Module):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
"""
def __init__(self, drop_prob=None):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
def get_shape(tensor):
shape = tensor.shape
if torch.onnx.is_in_onnx_export():
shape = [i.cpu().numpy() for i in shape]
return shape
class Conv2d_BN(nn.Sequential):
def __init__(self, a, b, ks=1, stride=1, pad=0, dilation=1,
groups=1, bn_weight_init=1, bias=False,
norm_cfg=dict(type='BN', requires_grad=True)):
super().__init__()
self.inp_channel = a
self.out_channel = b
self.ks = ks
self.pad = pad
self.stride = stride
self.dilation = dilation
self.groups = groups
# self.bias = bias
self.add_module('c', nn.Conv2d(
a, b, ks, stride, pad, dilation, groups, bias=bias))
bn = build_norm_layer(norm_cfg, b)[1]
nn.init.constant_(bn.weight, bn_weight_init)
nn.init.constant_(bn.bias, 0)
self.add_module('bn', bn)
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU, drop=0.,
norm_cfg=dict(type='BN', requires_grad=True)):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = Conv2d_BN(in_features, hidden_features, norm_cfg=norm_cfg)
self.dwconv = nn.Conv2d(hidden_features, hidden_features, 3, 1, 1, bias=True, groups=hidden_features)
self.act = act_layer()
self.fc2 = Conv2d_BN(hidden_features, out_features, norm_cfg=norm_cfg)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.dwconv(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class InvertedResidual(nn.Module):
def __init__(
self,
inp: int,
oup: int,
ks: int,
stride: int,
expand_ratio: int,
activations=None,
norm_cfg=dict(type='BN', requires_grad=True)
) -> None:
super(InvertedResidual, self).__init__()
self.stride = stride
self.expand_ratio = expand_ratio
assert stride in [1, 2]
if activations is None:
activations = nn.ReLU
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
layers = []
if expand_ratio != 1:
# pw
layers.append(Conv2d_BN(inp, hidden_dim, ks=1, norm_cfg=norm_cfg))
layers.append(activations())
layers.extend([
# dw
Conv2d_BN(hidden_dim, hidden_dim, ks=ks, stride=stride, pad=ks // 2, groups=hidden_dim, norm_cfg=norm_cfg),
activations(),
# pw-linear
Conv2d_BN(hidden_dim, oup, ks=1, norm_cfg=norm_cfg)
])
self.conv = nn.Sequential(*layers)
self.out_channels = oup
self._is_cn = stride > 1
def forward(self, x):
if self.use_res_connect:
return x + self.conv(x)
else:
return self.conv(x)
class StackedMV2Block(nn.Module):
def __init__(
self,
cfgs,
stem,
inp_channel=16,
activation=nn.ReLU,
norm_cfg=dict(type='BN', requires_grad=True),
width_mult=1.):
super().__init__()
self.stem = stem
if stem:
self.stem_block = nn.Sequential(
Conv2d_BN(3, inp_channel, 3, 2, 1, norm_cfg=norm_cfg),
activation()
)
self.cfgs = cfgs
self.layers = []
for i, (k, t, c, s) in enumerate(cfgs):
output_channel = _make_divisible(c * width_mult, 8)
exp_size = t * inp_channel
exp_size = _make_divisible(exp_size * width_mult, 8)
layer_name = 'layer{}'.format(i + 1)
layer = InvertedResidual(inp_channel, output_channel, ks=k, stride=s, expand_ratio=t, norm_cfg=norm_cfg,
activations=activation)
self.add_module(layer_name, layer)
inp_channel = output_channel
self.layers.append(layer_name)
def forward(self, x):
if self.stem:
x = self.stem_block(x)
for i, layer_name in enumerate(self.layers):
layer = getattr(self, layer_name)
x = layer(x)
return x
class SqueezeAxialPositionalEmbedding(nn.Module):
def __init__(self, dim, shape):
super().__init__()
self.pos_embed = nn.Parameter(torch.randn([1, dim, shape]))
def forward(self, x):
B, C, N = x.shape
x = x + F.interpolate(self.pos_embed, size=(N), mode='linear', align_corners=False)
return x
class Sea_Attention(torch.nn.Module):
def __init__(self, dim, key_dim=16, num_heads=4,
attn_ratio=2,
activation=None,
norm_cfg=dict(type='BN', requires_grad=True), ):
super().__init__()
self.num_heads = num_heads
self.scale = key_dim ** -0.5
self.key_dim = key_dim
self.nh_kd = nh_kd = key_dim * num_heads # num_head key_dim
self.d = int(attn_ratio * key_dim)
self.dh = int(attn_ratio * key_dim) * num_heads
self.attn_ratio = attn_ratio
self.to_q = Conv2d_BN(dim, nh_kd, 1, norm_cfg=norm_cfg)
self.to_k = Conv2d_BN(dim, nh_kd, 1, norm_cfg=norm_cfg)
self.to_v = Conv2d_BN(dim, self.dh, 1, norm_cfg=norm_cfg)
self.proj = torch.nn.Sequential(activation(), Conv2d_BN(
self.dh, dim, bn_weight_init=0, norm_cfg=norm_cfg))
self.proj_encode_row = torch.nn.Sequential(activation(), Conv2d_BN(
self.dh, self.dh, bn_weight_init=0, norm_cfg=norm_cfg))
self.pos_emb_rowq = SqueezeAxialPositionalEmbedding(nh_kd, 16)
self.pos_emb_rowk = SqueezeAxialPositionalEmbedding(nh_kd, 16)
self.proj_encode_column = torch.nn.Sequential(activation(), Conv2d_BN(
self.dh, self.dh, bn_weight_init=0, norm_cfg=norm_cfg))
self.pos_emb_columnq = SqueezeAxialPositionalEmbedding(nh_kd, 16)
self.pos_emb_columnk = SqueezeAxialPositionalEmbedding(nh_kd, 16)
self.dwconv = Conv2d_BN(self.dh + 2 * self.nh_kd, 2 * self.nh_kd + self.dh, ks=3, stride=1, pad=1, dilation=1,
groups=2 * self.nh_kd + self.dh, norm_cfg=norm_cfg)
self.act = activation()
self.pwconv = Conv2d_BN(2 * self.nh_kd + self.dh, dim, ks=1, norm_cfg=norm_cfg)
self.sigmoid = h_sigmoid()
def forward(self, x):
B, C, H, W = x.shape
q = self.to_q(x)
k = self.to_k(x)
v = self.to_v(x)
# detail enhance
qkv = torch.cat([q, k, v], dim=1)
qkv = self.act(self.dwconv(qkv))
qkv = self.pwconv(qkv)
# squeeze axial attention
## squeeze row
qrow = self.pos_emb_rowq(q.mean(-1)).reshape(B, self.num_heads, -1, H).permute(0, 1, 3, 2)
krow = self.pos_emb_rowk(k.mean(-1)).reshape(B, self.num_heads, -1, H)
vrow = v.mean(-1).reshape(B, self.num_heads, -1, H).permute(0, 1, 3, 2)
attn_row = torch.matmul(qrow, krow) * self.scale
attn_row = attn_row.softmax(dim=-1)
xx_row = torch.matmul(attn_row, vrow) # B nH H C
xx_row = self.proj_encode_row(xx_row.permute(0, 1, 3, 2).reshape(B, self.dh, H, 1))
## squeeze column
qcolumn = self.pos_emb_columnq(q.mean(-2)).reshape(B, self.num_heads, -1, W).permute(0, 1, 3, 2)
kcolumn = self.pos_emb_columnk(k.mean(-2)).reshape(B, self.num_heads, -1, W)
vcolumn = v.mean(-2).reshape(B, self.num_heads, -1, W).permute(0, 1, 3, 2)
attn_column = torch.matmul(qcolumn, kcolumn) * self.scale
attn_column = attn_column.softmax(dim=-1)
xx_column = torch.matmul(attn_column, vcolumn) # B nH W C
xx_column = self.proj_encode_column(xx_column.permute(0, 1, 3, 2).reshape(B, self.dh, 1, W))
xx = xx_row.add(xx_column)
xx = v.add(xx)
xx = self.proj(xx)
xx = self.sigmoid(xx) * qkv
return xx
class Sea_AttentionBlock(nn.Module):
def __init__(self, dim, key_dim=64, num_heads=4, mlp_ratio=2., attn_ratio=2., drop=0.,
drop_path=0.1, act_layer=nn.ReLU, norm_cfg=dict(type='BN2d', requires_grad=True)):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.mlp_ratio = mlp_ratio
self.attn = Sea_Attention(dim, key_dim=key_dim, num_heads=num_heads, attn_ratio=attn_ratio,
activation=act_layer, norm_cfg=norm_cfg)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop, norm_cfg=norm_cfg)
def forward(self, x1):
x1 = x1 + self.drop_path(self.attn(x1))
x1 = x1 + self.drop_path(self.mlp(x1))
return x1
class C2f_SeaformerBlock(nn.Module):
# CSP Bottleneck with 2 convolutions
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__()
self.c = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2)
self.m = nn.ModuleList(Sea_AttentionBlock(self.c) for _ in range(n))
def forward(self, x):
y = list(self.cv1(x).chunk(2, 1))
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
def forward_split(self, x):
y = list(self.cv1(x).split((self.c, self.c), 1))
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
class SeaformerBasicLayer(nn.Module):
def __init__(self, block_num, embedding_dim, key_dim, num_heads,
mlp_ratio=4., attn_ratio=2., drop=0., attn_drop=0., drop_path=0.,
norm_cfg=dict(type='BN2d', requires_grad=True),
act_layer=None):
super().__init__()
self.block_num = block_num
self.transformer_blocks = nn.ModuleList()
for i in range(self.block_num):
self.transformer_blocks.append(Sea_AttentionBlock(
embedding_dim, key_dim=key_dim, num_heads=num_heads,
mlp_ratio=mlp_ratio, attn_ratio=attn_ratio,
drop=drop, drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
norm_cfg=norm_cfg,
act_layer=act_layer))
def forward(self, x):
# token * N
for i in range(self.block_num):
x = self.transformer_blocks[i](x)
return x
class h_sigmoid(nn.Module):
def __init__(self, inplace=True):
super(h_sigmoid, self).__init__()
self.relu = nn.ReLU6(inplace=inplace)
def forward(self, x):
return self.relu(x + 3) / 6
def autopad(k, p=None, d=1): # kernel, padding, dilation
"""Pad to 'same' shape outputs."""
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):
"""Initialize Conv layer with given arguments including 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):
"""Apply convolution, batch normalization and activation to input tensor."""
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
"""Perform transposed convolution of 2D data."""
return self.act(self.conv(x))
class HGBlock_Sea(nn.Module):
"""
HG_Block of PPHGNetV2 with 2 convolutions and LightConv.
https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
"""
def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
"""Initializes a CSP Bottleneck with 1 convolution using specified input and output channels."""
super().__init__()
block = LightConv if lightconv else Conv
self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act) # squeeze conv
self.ec = Conv(c2 // 2, c2, 1, 1, act=act) # excitation conv
self.add = shortcut and c1 == c2
self.cv = Sea_AttentionBlock(c2)
def forward(self, x):
"""Forward pass of a PPHGNetV2 backbone layer."""
y = [x]
y.extend(m(y[-1]) for m in self.m)
y = self.cv(self.ec(self.sc(torch.cat(y, 1))))
return y + x if self.add else y
四、创新模块
4.1 改进点1⭐
模块改进方法
：直接加入
Sea_AttentionBlock模块
（
第五节讲解添加步骤
）。
Sea_AttentionBlock模块
添加后如下：
4.2 改进点2⭐
模块改进方法
：基于
Sea_AttentionBlock模块
的
HGBlock
（
第五节讲解添加步骤
）。
第二种改进方法是对
RT-DETR
中的
HGBlock模块
进行改进，并将
Sea_AttentionBlock
在加入到
HGBlock
模块中。
改进代码如下：
在
HGBlock
模块加入
Sea_AttentionBlock模块
，并将
HGBlock
重命名为
HGBlock_Sea
class HGBlock_Sea(nn.Module):
"""
HG_Block of PPHGNetV2 with 2 convolutions and LightConv.
https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
"""
def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
"""Initializes a CSP Bottleneck with 1 convolution using specified input and output channels."""
super().__init__()
block = LightConv if lightconv else Conv
self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act) # squeeze conv
self.ec = Conv(c2 // 2, c2, 1, 1, act=act) # excitation conv
self.add = shortcut and c1 == c2
self.cv = Sea_AttentionBlock(c2)
def forward(self, x):
"""Forward pass of a PPHGNetV2 backbone layer."""
y = [x]
y.extend(m(y[-1]) for m in self.m)
y = self.cv(self.ec(self.sc(torch.cat(y, 1))))
return y + x if self.add else y
注意❗：在
第五小节
中需要声明的模块名称为：
HGBlock_Sea
。
五、添加步骤
5.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
Sea_Attention.py
，将
第三节
中的代码粘贴到此处
5.2 修改二
在
AddModules
文件夹下新建
__init__.py
（已有则不用新建），在文件内导入模块：
from .Sea_Attention import *
5.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中，需要在两处位置添加各模块类名称。
首先：导入模块
其次：在
parse_model函数
中注册
HGBlock_Sea
模块
再在此函数中添加如下代码：
elif m in {Sea_AttentionBlock}:
args = [ch[f], *args]
六、yaml模型文件
6.1 模型改进版本1
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-Sea_AttentionBlock.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-Sea_AttentionBlock.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是在
骨干网络中
添加
Sea_AttentionBlock模块
。
# 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, 1, Sea_AttentionBlock, [1024]] # stage 4
- [-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, 18], 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, 13], 1, Concat, [1]] # cat Y5
- [-1, 3, RepC3, [256]] # F5 (27), pan_blocks.1
- [[22, 25, 28], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5)
6.2 模型改进版本2⭐
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-HGBlock_Sea.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-HGBlock_Sea.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
骨干网络
中的
HGBlock模块
替换成
HGBlock_Sea模块
。
# 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_Sea, [192, 1024, 5, True, False]] # cm, c2, k, light, shortcut
- [-1, 6, HGBlock_Sea, [192, 1024, 5, True, True]]
- [-1, 6, HGBlock_Sea, [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, 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)
七、成功运行结果
打印网络模型可以看到
Sea_AttentionBlock
和
HGBlock_Sea
已经加入到模型中，并可以进行训练了。
rtdetr-l-Sea_AttentionBlock ：
rtdetr-l-Sea_AttentionBlock summary: 731 layers, 213,701,827 parameters, 213,701,827 gradients, 150.5 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 1 180893696 ultralytics.nn.AddModules.Sea_Attention.Sea_AttentionBlock[1024, 1024]
10 -1 6 6708480 ultralytics.nn.modules.block.HGBlock [1024, 384, 2048, 5, 6, True, False]
11 -1 1 524800 ultralytics.nn.modules.conv.Conv [2048, 256, 1, 1, None, 1, 1, False]
12 -1 1 789760 ultralytics.nn.modules.transformer.AIFI [256, 1024, 8]
13 -1 1 66048 ultralytics.nn.modules.conv.Conv [256, 256, 1, 1]
14 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
15 7 1 262656 ultralytics.nn.modules.conv.Conv [1024, 256, 1, 1, None, 1, 1, False]
16 [-2, -1] 1 0 ultralytics.nn.modules.conv.Concat [1]
17 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
18 -1 1 66048 ultralytics.nn.modules.conv.Conv [256, 256, 1, 1]
19 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
20 3 1 131584 ultralytics.nn.modules.conv.Conv [512, 256, 1, 1, None, 1, 1, False]
21 [-2, -1] 1 0 ultralytics.nn.modules.conv.Concat [1]
22 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
23 -1 1 590336 ultralytics.nn.modules.conv.Conv [256, 256, 3, 2]
24 [-1, 18] 1 0 ultralytics.nn.modules.conv.Concat [1]
25 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
26 -1 1 590336 ultralytics.nn.modules.conv.Conv [256, 256, 3, 2]
27 [-1, 13] 1 0 ultralytics.nn.modules.conv.Concat [1]
28 -1 3 2232320 ultralytics.nn.modules.block.RepC3 [512, 256, 3]
29 [22, 25, 28] 1 7303907 ultralytics.nn.modules.head.RTDETRDecoder [1, [256, 256, 256]]
rtdetr-l-Sea_AttentionBlock summary: 731 layers, 213,701,827 parameters, 213,701,827 gradients, 150.5 GFLOPs
rtdetr-l-HGBlock_Sea ：
rtdetr-l-HGBlock_Sea summary: 832 layers, 55,015,619 parameters, 55,015,619 gradients, 174.1 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 9097856 ultralytics.nn.AddModules.Sea_Attention.HGBlock_Sea[512, 192, 1024, 5, 6, True, False]
6 -1 6 9458304 ultralytics.nn.AddModules.Sea_Attention.HGBlock_Sea[1024, 192, 1024, 5, 6, True, True]
7 -1 6 9458304 ultralytics.nn.AddModules.Sea_Attention.HGBlock_Sea[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 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 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 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-HGBlock_Sea summary: 832 layers, 55,015,619 parameters, 55,015,619 gradients, 174.1 GFLOPs