RT-DETR改进策略【Neck】| 有效且轻量的动态上采样算子：DySample
一、本文介绍
本文记录的是
利用DySample上采样对RT-DETR的颈部网络进行改进的方法研究
。
RT-DETR
采用传统的
最近邻插值
的方法进行上采样可能无法有效地捕捉特征的细节和语义信息，从而影响模型在密集预测任务中的性能。
DySample
通过
动态采样
的方式进行上采样，
能够更好地处理特征的细节和语义信息。
二、DySample介绍
通过学习采样来学习上采样
DySample
是一种超轻量且有效的动态上采样器，其设计出发点、原理和优势如下：
2.1、DySample原理
2.1.1 初步设计
通过PyTorch的内置函数，假设输入特征通过双线性插值被插值为连续的特征图，然后通过生成内容感知的采样点来重新采样该连续图。
具体实现为，给定特征图 X X X 和上采样尺度因子 s s s ，使用线性层生成偏移 O O O ，并通过Pixel Shuffling将其重塑为 2 × s H × s W 2 \times sH \times sW 2 × sH × s W ，然后将偏移 O O O 与原始采样网格 G G G 相加得到采样集 S S S ，最后通过网格采样函数根据采样集生成上采样后的特征图 X ′ X' X ′ 。
2.1.2 改进步骤
- 初始采样位置 ：在初步版本中，初始采样位置固定且分布不均匀，类似于“最近邻初始化”。为解决此问题，改为“双线性初始化”，即改变初始位置，使零偏移时能得到双线性插值的特征图，从而提高性能。
- 偏移范围 ：由于归一化层的存在，输出特征值的范围通常在 [ − 1 , 1 ] [ - 1, 1] [ − 1 , 1 ] ，导致局部采样位置的偏移范围可能重叠，影响边界预测并导致输出伪影。通过将偏移乘以0.25的“静态范围因子”，局部约束了采样位置的偏移范围，缓解了该问题。
- 分组 ：组向上采样，将特征图沿通道维度划分为 g g g 组，并为每组生成偏移。当 g = 4 g = 4 g = 4 时，性能得到提升。
- 动态范围因子 ：为增加偏移的灵活性，通过线性投影输入特征生成点级的“动态范围因子”，动态范围因子的值在 [ 0 , 0.5 ] [0, 0.5] [ 0 , 0.5 ] 范围内，以0.25为中心，进一步提升了性能。
- 偏移生成方式 ：研究了两种偏移生成方式，“线性 + 像素洗牌”（LP）和“像素洗牌 + 线性”（PL）。通过实验，根据不同模型设置了不同的组数量，并且发现PL版本在某些模型上表现更好，但在其他模型上略逊于LP版本。
2.1.3 最终变体
根据范围因子（静态/动态）和偏移生成方式（LP/PL），研究了四个变体：DySample（LP风格，静态范围因子）、DySample +（LP风格，动态范围因子）、DySample - S（PL风格，静态范围因子）、DySample - S +（PL风格，动态范围因子）。
2.2、优势
-
轻量高效
：与其他动态上采样器相比，
DySample不需要高分辨率的引导特征作为输入，也不需要除PyTorch之外的任何额外CUDA包，具有更少的推理延迟、内存占用、FLOPs和参数数量。 -
性能优越
：在五个密集预测任务（语义分割、目标检测、实例分割、全景分割和单目深度估计）中，与其他上采样器相比，
DySample报告了更好的性能。
论文： https://arxiv.org/pdf/2308.15085
源码： https://github.com/tiny-smart/dysample
三、DySample的实现代码
DySample模块
的实现代码如下：
import torch
import torch.nn as nn
import torch.nn.functional as F
def normal_init(module, mean=0, std=1, bias=0):
if hasattr(module, 'weight') and module.weight is not None:
nn.init.normal_(module.weight, mean, std)
if hasattr(module, 'bias') and module.bias is not None:
nn.init.constant_(module.bias, bias)
def constant_init(module, val, bias=0):
if hasattr(module, 'weight') and module.weight is not None:
nn.init.constant_(module.weight, val)
if hasattr(module, 'bias') and module.bias is not None:
nn.init.constant_(module.bias, bias)
class DySample(nn.Module):
def __init__(self, in_channels, scale=2, style='lp', groups=4, dyscope=False):
super().__init__()
self.scale = scale
self.style = style
self.groups = groups
assert style in ['lp', 'pl']
if style == 'pl':
assert in_channels >= scale ** 2 and in_channels % scale ** 2 == 0
assert in_channels >= groups and in_channels % groups == 0
if style == 'pl':
in_channels = in_channels // scale ** 2
out_channels = 2 * groups
else:
out_channels = 2 * groups * scale ** 2
self.offset = nn.Conv2d(in_channels, out_channels, 1)
normal_init(self.offset, std=0.001)
if dyscope:
self.scope = nn.Conv2d(in_channels, out_channels, 1, bias=False)
constant_init(self.scope, val=0.)
self.register_buffer('init_pos', self._init_pos())
def _init_pos(self):
h = torch.arange((-self.scale + 1) / 2, (self.scale - 1) / 2 + 1) / self.scale
return torch.stack(torch.meshgrid([h, h])).transpose(1, 2).repeat(1, self.groups, 1).reshape(1, -1, 1, 1)
def sample(self, x, offset):
B, _, H, W = offset.shape
offset = offset.view(B, 2, -1, H, W)
coords_h = torch.arange(H) + 0.5
coords_w = torch.arange(W) + 0.5
coords = torch.stack(torch.meshgrid([coords_w, coords_h])
).transpose(1, 2).unsqueeze(1).unsqueeze(0).type(x.dtype).to(x.device)
normalizer = torch.tensor([W, H], dtype=x.dtype, device=x.device).view(1, 2, 1, 1, 1)
coords = 2 * (coords + offset) / normalizer - 1
coords = F.pixel_shuffle(coords.view(B, -1, H, W), self.scale).view(
B, 2, -1, self.scale * H, self.scale * W).permute(0, 2, 3, 4, 1).contiguous().flatten(0, 1)
return F.grid_sample(x.reshape(B * self.groups, -1, H, W), coords, mode='bilinear',
align_corners=False, padding_mode="border").view(B, -1, self.scale * H, self.scale * W)
def forward_lp(self, x):
if hasattr(self, 'scope'):
offset = self.offset(x) * self.scope(x).sigmoid() * 0.5 + self.init_pos
else:
offset = self.offset(x) * 0.25 + self.init_pos
return self.sample(x, offset)
def forward_pl(self, x):
x_ = F.pixel_shuffle(x, self.scale)
if hasattr(self, 'scope'):
offset = F.pixel_unshuffle(self.offset(x_) * self.scope(x_).sigmoid(), self.scale) * 0.5 + self.init_pos
else:
offset = F.pixel_unshuffle(self.offset(x_), self.scale) * 0.25 + self.init_pos
return self.sample(x, offset)
def forward(self, x):
if self.style == 'pl':
return self.forward_pl(x)
return self.forward_lp(x)
四、添加步骤
4.1 修改一
① 在
ultralytics/nn/
目录下新建
AddModules
文件夹用于存放模块代码
② 在
AddModules
文件夹下新建
DySample.py
，将
第三节
中的代码粘贴到此处
5.2 修改二
在
AddModules
文件夹下新建
__init__.py
（已有则不用新建），在文件内导入模块：
from .DySample import *
5.3 修改三
在
ultralytics/nn/modules/tasks.py
文件中，需要在两处位置添加各模块类名称。
首先：导入模块
然后，在
parse_model函数
中添加如下代码：
elif m is DySample:
args = [ch[f]]
五、yaml模型文件
5.1 模型改进版本⭐
此处以
ultralytics/cfg/models/rt-detr/rtdetr-l.yaml
为例，在同目录下创建一个用于自己数据集训练的模型文件
rtdetr-l-DySample.yaml
。
将
rtdetr-l.yaml
中的内容复制到
rtdetr-l-DySample.yaml
文件下，修改
nc
数量等于自己数据中目标的数量。
📌 模型的修改方法是将
颈部网络
中的
上采样
替换成
DySample模块
。
# 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, DySample, []]
- [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, DySample, []]
- [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)
六、成功运行结果
打印网络模型可以看到
DySample
已经加入到模型中，并可以进行训练了。
rtdetr-l-DySample ：
rtdetr-l-DySample summary: 683 layers, 32,824,579 parameters, 32,824,579 gradients, 108.0 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 8224 ultralytics.nn.AddModules.DySample.DySample [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 8224 ultralytics.nn.AddModules.DySample.DySample [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-DySample summary: 683 layers, 32,824,579 parameters, 32,824,579 gradients, 108.0 GFLOPs