RT-DETR改进策略【卷积层】| CVPR-2024 PKI Module 获取多尺度纹理特征，适应尺度变化大的目标
一、本文介绍
本文记录的是
利用
PKI Module
优化
RT-DETR
的目标检测方法研究
。在
遥感图像
目标检测中，与一般目标检测不同，需要在单张图像中定位和识别不同大小的目标。为了解决目标尺度变化大的挑战，本文引入
PKI Module
来捕获多尺度纹理特征，实验验证，有效涨点。
二、PKI Module原理介绍
Poly Kernel Inception Network for Remote Sensing Detection
PKI Module
是
Poly Kernel Inception Network (PKINet)
中的一个重要模块，其设计原理、结构和优势如下：
2.1 原理
PKI Module
是一个
Inception-Style
模块，
通过不同尺寸的卷积核组合来提取不同尺度的特征
。它先使用小卷积核卷积抓取局部信息，然后使用一组并行的深度可分离卷积来捕获多尺度的上下文信息。通过这种方式，
可以在不同感受野上提取特征，并将局部和上下文特征进行融合，以获取更丰富的特征表示，同时避免因单一尺度卷积核或扩张卷积带来的问题
，如小卷积核可能丢失长距离上下文信息，大卷积核可能引入背景噪声或生成过于稀疏的特征表示。
2.2 结构
2.2.1 局部特征提取
- 对于第 l l l 阶段第 n n n 个PKI Block中的PKI Module，首先通过 k s × k s k_{s}×k_{s} k s × k s 卷积（在实验中 k s = 3 k_{s}=3 k s = 3 ）对输入 X l − 1 , n ( 2 ) X_{l - 1,n}^{(2)} X l − 1 , n ( 2 ) 进行局部特征提取，得到 L l − 1 , n ∈ R 1 2 C l × H l × W l L_{l - 1,n}\in\mathbb{R}^{\frac{1}{2}C_{l}×H_{l}×W_{l}} L l − 1 , n ∈ R 2 1 C l × H l × W l 。
2.2.2 多尺度上下文特征提取
- 接着使用一组并行的深度可分离卷积 D W C o n v k ( m ) × k ( m ) DWConv_{k^{(m)}×k^{(m)}} D W C o n v k ( m ) × k ( m ) （ m = 1 , ⋯ , 4 m = 1,\cdots,4 m = 1 , ⋯ , 4 ，其中 k ( m ) = ( m + 1 ) × 2 + 1 k^{(m)}=(m + 1)×2+1 k ( m ) = ( m + 1 ) × 2 + 1 ）对 L l − 1 , n L_{l - 1,n} L l − 1 , n 进行卷积，得到不同尺度的上下文特征 Z l − 1 , n ( m ) ∈ R 1 2 C l × H l × W l Z_{l - 1,n}^{(m)}\in\mathbb{R}^{\frac{1}{2}C_{l}×H_{l}×W_{l}} Z l − 1 , n ( m ) ∈ R 2 1 C l × H l × W l 。
2.2.3 特征融合
- 将局部特征 L l − 1 , n L_{l - 1,n} L l − 1 , n 和多尺度上下文特征 ∑ m = 1 4 Z l − 1 , n ( m ) \sum_{m = 1}^{4}Z_{l - 1,n}^{(m)} ∑ m = 1 4 Z l − 1 , n ( m ) 相加后，通过一个 1 × 1 1×1 1 × 1 卷积进行融合，得到输出特征 P l − 1 , n ∈ R 1 2 C i × H i × W i P_{l - 1,n}\in\mathbb{R}^{\frac{1}{2}C_{i}×H_{i}×W_{i}} P l − 1 , n ∈ R 2 1 C i × H i × W i 。这个 1 × 1 1×1 1 × 1 卷积起到了通道融合的作用，能够整合不同感受野大小的特征。
2.3 优势
-
多尺度特征提取
- 能够有效捕获多尺度纹理特征，适应遥感图像中目标尺度变化大的特点，通过不同大小的卷积核组合，可以获取不同尺度的局部和上下文信息，提高对不同大小目标的检测能力。
-
避免特征稀疏和噪声问题
- 不使用扩张卷积，防止了提取过于稀疏的特征表示，同时相比于单纯使用大卷积核，避免了引入过多背景噪声，从而提高了特征的质量和检测的准确性。
-
特征融合优势
- 通过 1 × 1 1×1 1 × 1 卷积进行特征融合，能够合理地整合不同尺度的特征，在保留局部纹理特征完整性的同时，捕获到更广泛的上下文信息，使得提取的特征更具代表性和判别力。
论文： https://arxiv.org/pdf/2403.06258
源码： https://github.com/NUST-Machine-Intelligence-Laboratory/PKINet
三、PKI Module的实现代码
InceptionBottleneck模块
的实现代码如下：
from typing import Optional, Sequence
import numpy as np
import torch.nn as nn
from mmengine.model import BaseModule
import torch
class CAA(BaseModule):
"""Context Anchor Attention"""
def __init__(
self,
channels: int,
h_kernel_size: int = 11,
v_kernel_size: int = 11,
norm_cfg: Optional[dict] = dict(type='BN', momentum=0.03, eps=0.001),
act_cfg: Optional[dict] = dict(type='SiLU'),
init_cfg: Optional[dict] = None,
):
super().__init__(init_cfg)
self.avg_pool = nn.AvgPool2d(7, 1, 3)
self.conv1 = ConvModule(channels, channels, 1, 1, 0,
norm_cfg=norm_cfg, act_cfg=act_cfg)
self.h_conv = ConvModule(channels, channels, (1, h_kernel_size), 1,
(0, h_kernel_size // 2), groups=channels,
norm_cfg=None, act_cfg=None)
self.v_conv = ConvModule(channels, channels, (v_kernel_size, 1), 1,
(v_kernel_size // 2, 0), groups=channels,
norm_cfg=None, act_cfg=None)
self.conv2 = ConvModule(channels, channels, 1, 1, 0,
norm_cfg=norm_cfg, act_cfg=act_cfg)
self.act = nn.Sigmoid()
def forward(self, x):
attn_factor = self.act(self.conv2(self.v_conv(self.h_conv(self.conv1(self.avg_pool(x))))))
return attn_factor
def autopad(kernel_size: int, padding: Optional[int] = None, dilation: int = 1) -> int:
"""Calculate the padding size based on kernel size and dilation."""
if padding is None:
padding = (kernel_size - 1) * dilation // 2
return padding
def make_divisible(value: int, divisor: int = 8) -> int:
"""Make a value divisible by a certain divisor."""
return int((value + divisor // 2) // divisor * divisor)
class ConvModule(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
norm_cfg: Optional[dict] = None,
act_cfg: Optional[dict] = None):
super().__init__()
layers = []
layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation=dilation, groups=groups, bias=(norm_cfg is None)))
if norm_cfg:
norm_layer = self._get_norm_layer(out_channels, norm_cfg)
layers.append(norm_layer)
if act_cfg:
act_layer = self._get_act_layer(act_cfg)
layers.append(act_layer)
self.block = nn.Sequential(*layers)
def forward(self, x):
return self.block(x)
def _get_norm_layer(self, num_features, norm_cfg):
if norm_cfg['type'] == 'BN':
return nn.BatchNorm2d(num_features, momentum=norm_cfg.get('momentum', 0.1), eps=norm_cfg.get('eps', 1e-5))
# Add more normalization types if needed
raise NotImplementedError(f"Normalization layer '{norm_cfg['type']}' is not implemented.")
def _get_act_layer(self, act_cfg):
if act_cfg['type'] == 'ReLU':
return nn.ReLU(inplace=True)
if act_cfg['type'] == 'SiLU':
return nn.SiLU(inplace=True)
# Add more activation types if needed
raise NotImplementedError(f"Activation layer '{act_cfg['type']}' is not implemented.")
# Update InceptionBottleneck's constructor call to avoid conflicts
class InceptionBottleneck(nn.Module):
"""Bottleneck with Inception module"""
def __init__(
self,
in_channels: int,
out_channels: Optional[int] = None,
kernel_sizes: Sequence[int] = (3, 5, 7, 9, 11),
dilations: Sequence[int] = (1, 1, 1, 1, 1),
expansion: float = 1.0,
add_identity: bool = True,
with_caa: bool = True,
caa_kernel_size: int = 11,
norm_cfg: Optional[dict] = dict(type='BN', momentum=0.03, eps=0.001),
act_cfg: Optional[dict] = dict(type='SiLU')):
super().__init__()
out_channels = out_channels or in_channels
hidden_channels = make_divisible(int(out_channels * expansion), 8)
self.pre_conv = ConvModule(in_channels, hidden_channels, 1, 1, 0,
norm_cfg=norm_cfg, act_cfg=act_cfg)
self.dw_conv = ConvModule(hidden_channels, hidden_channels, kernel_sizes[0], 1,
autopad(kernel_sizes[0], None, dilations[0]),
dilation=dilations[0], groups=hidden_channels,
norm_cfg=None, act_cfg=None)
self.dw_conv1 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[1], 1,
autopad(kernel_sizes[1], None, dilations[1]),
dilation=dilations[1], groups=hidden_channels,
norm_cfg=None, act_cfg=None)
self.dw_conv2 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[2], 1,
autopad(kernel_sizes[2], None, dilations[2]),
dilation=dilations[2], groups=hidden_channels,
norm_cfg=None, act_cfg=None)
self.dw_conv3 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[3], 1,
autopad(kernel_sizes[3], None, dilations[3]),
dilation=dilations[3], groups=hidden_channels,
norm_cfg=None, act_cfg=None)
self.dw_conv4 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[4], 1,
autopad(kernel_sizes[4], None, dilations[4]),
dilation=dilations[4], groups=hidden_channels,
norm_cfg=None, act_cfg=None)
self.pw_conv = ConvModule(hidden_channels, hidden_channels, 1, 1, 0,
norm_cfg=norm_cfg, act_cfg=act_cfg)
if with_caa:
self.caa_factor = CAA(hidden_channels, caa_kernel_size, caa_kernel_size, None, None)
else:
self.caa_factor = None
self.add_identity = add_identity and in_channels == out_channels
self.post_conv = ConvModule(hidden_channels, out_channels, 1, 1, 0,
norm_cfg=norm_cfg, act_cfg=act_cfg)
def forward(self, x):
x = self.pre_conv(x)
y = x
x = self.dw_conv(x)
x = x + self.dw_conv1(x) + self.dw_conv2(x) + self.dw_conv3(x) + self.dw_conv4(x)
x = self.pw_conv(x)
if self.caa_factor is not None:
y = self.caa_factor(y)
if self.add_identity:
y = x * y
x = x + y
else:
x = x * y
x = self.post_conv(x)
return x
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 RepConv(nn.Module):
"""
RepConv is a basic rep-style block, including training and deploy status.
This module is used in RT-DETR.
Based on https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py
"""
default_act = nn.SiLU() # default activation
def __init__(self, c1, c2, k=3, s=1, p=1, g=1, d=1, act=True, bn=False, deploy=False):
"""Initializes Light Convolution layer with inputs, outputs & optional activation function."""
super().__init__()
assert k == 3 and p == 1
self.g = g
self.c1 = c1
self.c2 = c2
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
self.bn = nn.BatchNorm2d(num_features=c1) if bn and c2 == c1 and s == 1 else None
self.conv1 = Conv(c1, c2, k, s, p=p, g=g, act=False)
self.conv2 = Conv(c1, c2, 1, s, p=(p - k // 2), g=g, act=False)
def forward_fuse(self, x):
"""Forward process."""
return self.act(self.conv(x))
def forward(self, x):
"""Forward process."""
id_out = 0 if self.bn is None else self.bn(x)
return self.act(self.conv1(x) + self.conv2(x) + id_out)
def get_equivalent_kernel_bias(self):
"""Returns equivalent kernel and bias by adding 3x3 kernel, 1x1 kernel and identity kernel with their biases."""
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
kernelid, biasid = self._fuse_bn_tensor(self.bn)
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
@staticmethod
def _pad_1x1_to_3x3_tensor(kernel1x1):
"""Pads a 1x1 tensor to a 3x3 tensor."""
if kernel1x1 is None:
return 0
else:
return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])
def _fuse_bn_tensor(self, branch):
"""Generates appropriate kernels and biases for convolution by fusing branches of the neural network."""
if branch is None:
return 0, 0
if isinstance(branch, Conv):
kernel = branch.conv.weight
running_mean = branch.bn.running_mean
running_var = branch.bn.running_var
gamma = branch.bn.weight
beta = branch.bn.bias
eps = branch.bn.eps
elif isinstance(branch, nn.BatchNorm2d):
if not hasattr(self, "id_tensor"):
input_dim = self.c1 // self.g
kernel_value = np.zeros((self.c1, input_dim, 3, 3), dtype=np.float32)
for i in range(self.c1):
kernel_value[i, i % input_dim, 1, 1] = 1
self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
kernel = self.id_tensor
running_mean = branch.running_mean
running_var = branch.running_var
gamma = branch.weight
beta = branch.bias
eps = branch.eps
std = (running_var + eps).sqrt()
t = (gamma / std).reshape(-1, 1, 1, 1)
return kernel * t, beta - running_mean * gamma / std
def fuse_convs(self):
"""Combines two convolution layers into a single layer and removes unused attributes from the class."""
if hasattr(self, "conv"):
return
kernel, bias = self.get_equivalent_kernel_bias()
self.conv = nn.Conv2d(
in_channels=self.conv1.conv.in_channels,
out_channels=self.conv1.conv.out_channels,
kernel_size=self.conv1.conv.kernel_size,
stride=self.conv1.conv.stride,
padding=self.conv1.conv.padding,
dilation=self.conv1.conv.dilation,
groups=self.conv1.conv.groups,
bias=True,
).requires_grad_(False)
self.conv.weight.data = kernel
self.conv.bias.data = bias
for para in self.parameters():
para.detach_()
self.__delattr__("conv1")
self.__delattr__("conv2")
if hasattr(self, "nm"):
self.__delattr__("nm")
if hasattr(self, "bn"):
self.__delattr__("bn")
if hasattr(self, "id_tensor"):
self.__delattr__("id_tensor")
class RepC3_PKI(nn.Module):
"""Rep C3."""
def __init__(self, c1, c2, n=3, e=1.0):
"""Initialize CSP Bottleneck with a single convolution using input channels, output channels, and number."""
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c2, 1, 1)
self.cv2 = InceptionBottleneck(c1, c2)
self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()
def forward(self, x):
"""Forward pass of RT-DETR neck layer."""
return self.cv3(self.m(self.cv1(x)) + self.cv2(x))
四、创新模块
4.1 改进点⭐
模块改进方法
：直接加入
InceptionBottleneck
（
第五节讲解添加步骤
）。
InceptionBottleneck
模块加入如下：
4.2 改进点⭐
模块改进方法
：基于
InceptionBottleneck模块
的
RepC3
（
第五节讲解添加步骤
）。
第二种改进方法是对
rtdetr-l
中的
RepC3模块
进行改进，并将
InceptionBottleneck
在加入到
RepC3
模块中。
改进代码如下：
对
RepC3
模块进行改进，加入
InceptionBottleneck模块
class RepC3_PKI(nn.Module):
"""Rep C3."""
def __init__(self, c1, c2, n=3, e=1.0):
"""Initialize CSP Bottleneck with a single convolution using input channels, output channels, and number."""
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c2, 1, 1)
self.cv2 = InceptionBottleneck(c1, c2)
self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()
def forward(self, x):
"""Forward pass of RT-DETR neck layer."""
return self.cv3(self.m(self.cv1(x)) + self.cv2(x))
注意❗：在
第五小节
中需要声明的模块名称为：
InceptionBottleneck
和
RepC3_PKI
。
五、添加步骤
5.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
PKIModule.py
，将
第三节
中的代码粘贴到此处
5.2 修改二
在
AddModules
文件夹下新建
__init__.py
（已有则不用新建），在文件内导入模块：
from .PKIModule import *
5.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中，需要在指定位置添加各模块类名称。
首先：导入模块
其次：在
parse_model函数
中注册
InceptionBottleneck
和
RepC3_PKI
模块
六、yaml模型文件
6.1 模型改进版本⭐
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-PKIModule.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-PKIModule.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
骨干网络
中的
HGBlock模块
替换成
PKIModule模块
。
# 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, 1, InceptionBottleneck, [512]] # cm, c2, k, light, shortcut
- [-1, 1, InceptionBottleneck, [512]]
- [-1, 1, InceptionBottleneck, [512]] # 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)
6.2 模型改进版本⭐
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-RepC3_PKI.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-RepC3_PKI.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
颈部网络
中的
RepC3模块
替换成
RepC3_PKI模块
。
# 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, Concat, [1]]
- [-1, 1, RepC3_PKI, [512]] # 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, 1, RepC3_PKI, [512]] # X3 (21), fpn_blocks.1
- [-1, 1, Conv, [256, 3, 2]] # 22, downsample_convs.0
- [[-1, 17], 1, Concat, [1]] # cat Y4
- [-1, 1, RepC3_PKI, [512]] # F4 (24), pan_blocks.0
- [-1, 1, Conv, [256, 3, 2]] # 25, downsample_convs.1
- [[-1, 12], 1, Concat, [1]] # cat Y5
- [-1, 1, RepC3_PKI, [512]] # F5 (27), pan_blocks.1
- [[21, 24, 27], 1, RTDETRDecoder, [nc]] # Detect(P3, P4, P5)
七、成功运行结果
打印网络模型可以看到
InceptionBottleneck
和
C2fCIB_PKI
已经加入到模型中，并可以进行训练了。
rtdetr-l-PKIModule ：
rtdetr-l-PKIModule summary: 651 layers, 31,296,835 parameters, 31,296,835 gradients, 103.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 1 1475584 ultralytics.nn.AddModules.PKIModule.InceptionBottleneck[512, 512]
6 -1 1 1475584 ultralytics.nn.AddModules.PKIModule.InceptionBottleneck[512, 512]
7 -1 1 1475584 ultralytics.nn.AddModules.PKIModule.InceptionBottleneck[512, 512]
8 -1 1 11264 ultralytics.nn.modules.conv.DWConv [512, 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 131584 ultralytics.nn.modules.conv.Conv [512, 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-PKIModule summary: 651 layers, 31,296,835 parameters, 31,296,835 gradients, 103.1 GFLOPs
rtdetr-l-RepC3_PKI ：
rtdetr-l-RepC3_PKI summary: 854 layers, 63,757,507 parameters, 63,757,507 gradients, 259.2 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 0 ultralytics.nn.modules.conv.Concat [1]
16 -1 1 9609216 ultralytics.nn.AddModules.PKIModule.RepC3_PKI[512, 512]
17 -1 1 131584 ultralytics.nn.modules.conv.Conv [512, 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 1 9609216 ultralytics.nn.AddModules.PKIModule.RepC3_PKI[512, 512]
22 -1 1 1180160 ultralytics.nn.modules.conv.Conv [512, 256, 3, 2]
23 [-1, 17] 1 0 ultralytics.nn.modules.conv.Concat [1]
24 -1 1 9609216 ultralytics.nn.AddModules.PKIModule.RepC3_PKI[512, 512]
25 -1 1 1180160 ultralytics.nn.modules.conv.Conv [512, 256, 3, 2]
26 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]
27 -1 1 9609216 ultralytics.nn.AddModules.PKIModule.RepC3_PKI[512, 512]
28 [21, 24, 27] 1 7500515 ultralytics.nn.modules.head.RTDETRDecoder [1, [512, 512, 512]]
rtdetr-l-RepC3_PKI summary: 854 layers, 63,757,507 parameters, 63,757,507 gradients, 259.2 GFLOPs