大家好,我是Snu77,这里是RT-DETR有效涨点专栏。
本专栏的内容为根据ultralytics版本的RT-DETR进行改进,内容持续更新,每周更新文章数量3-10篇。
专栏以ResNet18、ResNet50为基础修改版本,同时修改内容也支持ResNet32、ResNet101和PPHGNet版本,其中ResNet为RT-DETR官方版本1:1移植过来的,参数量基本保持一致(误差很小很小),不同于ultralytics仓库版本的ResNet官方版本,同时ultralytics仓库的一些参数是和RT-DETR相冲的所以我也是会教大家调好一些参数和代码,真正意义上的跑ultralytics的和RT-DETR官方版本的无区别
👑欢迎大家订阅本专栏,一起学习RT-DETR👑??
本文给大家带来的改进机制是TransNeXt特征提取网络,其发表于2023年的12月份是一个最新最前沿的网络模型,将其应用在我们的特征提取网络来提取特征,同时本文给大家解决其自带的一个报错,通过结合聚合的像素聚焦注意力和卷积GLU,模拟生物视觉系统,特别是对于中心凹的视觉感知。这种方法使得每个像素都能实现全局感知,并强化了模型的信息混合和自然视觉感知能力。TransNeXt在各种视觉任务中,包括图像分类、目标检测和语义分割,都显示出优异的性能。
目录
官方论文地址:官方论文地址
官方代码地址:官方代码地址
TransNeXt: Robust Foveal Visual Perception for Vision Transformers介绍了一种新的视觉模型,旨在改进现有视觉变换器的性能。这个模型,被称为 TransNeXt,通过结合聚合的像素聚焦注意力和卷积GLU,模拟生物视觉系统,特别是对于中心凹的视觉感知。这种方法使得每个像素都能实现全局感知,并强化了模型的信息混合和自然视觉感知能力。TransNeXt在各种视觉任务中,包括图像分类、目标检测和语义分割,都显示出优异的性能。
TransNeXt的主要创新点包括:
1. 聚合注意力机制:模仿生物中心凹视觉和连续眼动,使每个令牌在特征图上都能实现全球感知。
2. 卷积GLU(Gated Linear Unit):弥补了GLU和SE(Squeeze-and-Excitation)机制之间的差距,增强局部建模能力和模型鲁棒性。
这些创新点共同使TransNeXt在图像分类、目标检测和语义分割等多种视觉任务中表现卓越。
这幅图展示了不同视觉信息聚合机制的比较,包括提出的方法和生物视觉系统。通过平均超过5000张2242分辨率的ImageNet-1K验证集图像,展示了不同方法与流行背景网络之间的有效感受野(ERF)的可视化对比。图中展示了四种不同的注意力机制:局部注意力、滑动窗口注意力与卷积、池化注意力,以及本文提出的像素聚焦注意力,以及它们与生物视觉系统的对比。每种机制下,红星表示当前查询的位置,黑色区域表示当前查询无法感知的区域。右侧的图表比较了TransNeXt(本文提出的方法)和其他几种流行的模型在处理视觉信息时的差异。
聚合注意力机制(Aggregated Attention, AA)是TransNeXt模型中的一个核心创新,它融合了多种注意力机制并为多尺度输入增强了外推能力,具体包括以下几点:
1. 像素聚焦注意力:
2. 集成多样的注意力机制:
3. 位置注意力:
4. 多尺度输入的外推能力:
5. 聚合注意力(Aggregated Attention):
总结:聚合注意力机制通过模拟生物视觉系统,提供了一种更自然的视觉感知方式,可以有效地处理来自不同层次和尺度的信息,并通过结合不同的注意力路径和可学习组件,增强了模型对于多尺度输入的外推能力。
缺点:但是它总结了这么多注意力机制,它限制了通道数减少了参数量,但是其运算非常复杂导致速度很慢
图3是像素聚焦注意力(左图)与聚合注意力(右图)之间对比的插图。两者都具有10x10的特征尺寸,一个3x3的窗口尺寸和2x2的池化尺寸。
左图(像素聚焦注意力):
右图(聚合注意力):
卷积GLU是TransNeXt模型中的一个关键创新点,旨在弥补GLU和SE(大家看的熟悉么我们之前讲过)机制之间的差距。以下是关于卷积GLU的详细介绍:
1. 基于最近邻图像特征的通道注意力:卷积GLU采用了一种基于最近邻图像特征的通道注意力机制。这种设计避免了SE机制中全局平均池化的过于粗粒度的缺点,并满足了一些没有位置编码设计的ViT(视觉变换器)模型的需求,这些模型需要通过深度卷积提供的位置信息。
2.强化局部建模能力和模型鲁棒性:与传统的卷积前馈网络相比,卷积GLU通过较少的浮点运算(FLOPs)实现了通道混合器的注意力化,从而有效地增强了模型的鲁棒性。
3. 创建新的视觉主干网络TransNeXt:将聚合注意力和卷积GLU结合起来,创造了一个新的视觉主干网络,名为TransNeXt。通过广泛的实验,TransNeXt在多个模型尺寸上都实现了最先进的性能。
总结:卷积GLU的引入,使得每个令牌都能够基于其最近的细粒度特征拥有独特的门控信号,这不仅提高了对局部特征的建模能力,还提高了模型在处理不同尺度和复杂性的视觉数据时的稳健性。
图4展示了当前流行的通道混合器设计与卷积GLU(Convolutional Gated Linear Unit)的比较。四个框架分别表示:
1. 原始前馈网络(Original Feed-Forward):
2. 卷积前馈网络(Convolutional Feed-Forward):
3. 门控线性单元(Gated Linear Unit, GLU):
4. 带有SE模块的前馈网络(FFN with SE module):
5. 卷积门控线性单元(Convolutional Gated Linear Unit):
图5展示了TransNeXt架构的一个示意图,揭示了其内部的组件和数据流。这个架构通过多个阶段的处理来处理输入图像,其中每个阶段都包含卷积GLU和聚合注意力机制的层。以下是每个阶段的详细介绍:
1. 图像输入:
2. 阶段1:
3. 阶段2和3:
4. 阶段4:
总结:通过这些阶段的处理,TransNeXt模型能够逐步提取和处理图像特征,从局部像素级特征到更高层次的抽象表示。每个阶段的输出都准备好进入下一个阶段,直到最终生成能够用于图像分类、目标检测或语义分割任务的高级特征。此架构展示了如何通过结合卷积和注意力机制来有效地处理视觉数据,同时逐步增加通道数和降低空间分辨率,以提高计算效率和模型性能。
官方的代码提供了两种一种是需要编译的,条件非常苛刻,但是速度运行比不编译的要快一些,我这里提供的是不编译的方便大家运行,大家要是想用编译的可以从我给的代码链接中找到,按照我的对比着来进行修改即可。
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from functools import partial
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
import math
class DWConv(nn.Module):
def __init__(self, dim=768):
super(DWConv, self).__init__()
self.dwconv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=True, groups=dim)
def forward(self, x, H, W):
B, N, C = x.shape
x = x.transpose(1, 2).view(B, C, H, W).contiguous()
x = self.dwconv(x)
x = x.flatten(2).transpose(1, 2)
return x
class ConvolutionalGLU(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
hidden_features = int(2 * hidden_features / 3)
self.fc1 = nn.Linear(in_features, hidden_features * 2)
self.dwconv = DWConv(hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x, H, W):
x, v = self.fc1(x).chunk(2, dim=-1)
x = self.act(self.dwconv(x, H, W)) * v
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
@torch.no_grad()
def get_relative_position_cpb(query_size, key_size, pretrain_size=None):
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
pretrain_size = pretrain_size or query_size
axis_qh = torch.arange(query_size[0], dtype=torch.float32)
axis_kh = F.adaptive_avg_pool1d(axis_qh.unsqueeze(0), key_size[0]).squeeze(0)
axis_qw = torch.arange(query_size[1], dtype=torch.float32)
axis_kw = F.adaptive_avg_pool1d(axis_qw.unsqueeze(0), key_size[1]).squeeze(0)
axis_kh, axis_kw = torch.meshgrid(axis_kh, axis_kw)
axis_qh, axis_qw = torch.meshgrid(axis_qh, axis_qw)
axis_kh = torch.reshape(axis_kh, [-1])
axis_kw = torch.reshape(axis_kw, [-1])
axis_qh = torch.reshape(axis_qh, [-1])
axis_qw = torch.reshape(axis_qw, [-1])
relative_h = (axis_qh[:, None] - axis_kh[None, :]) / (pretrain_size[0] - 1) * 8
relative_w = (axis_qw[:, None] - axis_kw[None, :]) / (pretrain_size[1] - 1) * 8
relative_hw = torch.stack([relative_h, relative_w], dim=-1).view(-1, 2)
relative_coords_table, idx_map = torch.unique(relative_hw, return_inverse=True, dim=0)
relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
torch.abs(relative_coords_table) + 1.0) / torch.log2(torch.tensor(8, dtype=torch.float32))
return idx_map, relative_coords_table
@torch.no_grad()
def get_seqlen_and_mask(input_resolution, window_size):
attn_map = F.unfold(torch.ones([1, 1, input_resolution[0], input_resolution[1]]), window_size,
dilation=1, padding=(window_size // 2, window_size // 2), stride=1)
attn_local_length = attn_map.sum(-2).squeeze().unsqueeze(-1)
attn_mask = (attn_map.squeeze(0).permute(1, 0)) == 0
return attn_local_length, attn_mask
class AggregatedAttention(nn.Module):
def __init__(self, dim, input_resolution, num_heads=8, window_size=3, qkv_bias=True,
attn_drop=0., proj_drop=0., sr_ratio=1):
super().__init__()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.sr_ratio = sr_ratio
assert window_size % 2 == 1, "window size must be odd"
self.window_size = window_size
self.local_len = window_size ** 2
self.pool_H, self.pool_W = input_resolution[0] // self.sr_ratio, input_resolution[1] // self.sr_ratio
self.pool_len = self.pool_H * self.pool_W
self.unfold = nn.Unfold(kernel_size=window_size, padding=window_size // 2, stride=1)
self.temperature = nn.Parameter(
torch.log((torch.ones(num_heads, 1, 1) / 0.24).exp() - 1)) # Initialize softplus(temperature) to 1/0.24.
self.q = nn.Linear(dim, dim, bias=qkv_bias)
self.query_embedding = nn.Parameter(
nn.init.trunc_normal_(torch.empty(self.num_heads, 1, self.head_dim), mean=0, std=0.02))
self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
# Components to generate pooled features.
self.pool = nn.AdaptiveAvgPool2d((self.pool_H, self.pool_W))
self.sr = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0)
self.norm = nn.LayerNorm(dim)
self.act = nn.GELU()
# mlp to generate continuous relative position bias
self.cpb_fc1 = nn.Linear(2, 512, bias=True)
self.cpb_act = nn.ReLU(inplace=True)
self.cpb_fc2 = nn.Linear(512, num_heads, bias=True)
# relative bias for local features
self.relative_pos_bias_local = nn.Parameter(
nn.init.trunc_normal_(torch.empty(num_heads, self.local_len), mean=0,
std=0.0004))
# Generate padding_mask && sequnce length scale
local_seq_length, padding_mask = get_seqlen_and_mask(input_resolution, window_size)
self.register_buffer("seq_length_scale", torch.as_tensor(np.log(local_seq_length.numpy() + self.pool_len)),
persistent=False)
self.register_buffer("padding_mask", padding_mask, persistent=False)
# dynamic_local_bias:
self.learnable_tokens = nn.Parameter(
nn.init.trunc_normal_(torch.empty(num_heads, self.head_dim, self.local_len), mean=0, std=0.02))
self.learnable_bias = nn.Parameter(torch.zeros(num_heads, 1, self.local_len))
def forward(self, x, H, W, relative_pos_index, relative_coords_table):
B, N, C = x.shape
# Generate queries, normalize them with L2, add query embedding, and then magnify with sequence length scale and temperature.
# Use softplus function ensuring that the temperature is not lower than 0.
q_norm = F.normalize(self.q(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3), dim=-1)
q_norm_scaled = (q_norm + self.query_embedding) * F.softplus(self.temperature) * self.seq_length_scale
# Generate unfolded keys and values and l2-normalize them
k_local, v_local = self.kv(x).chunk(2, dim=-1)
k_local = F.normalize(k_local.reshape(B, N, self.num_heads, self.head_dim), dim=-1).reshape(B, N, -1)
kv_local = torch.cat([k_local, v_local], dim=-1).permute(0, 2, 1).reshape(B, -1, H, W)
k_local, v_local = self.unfold(kv_local).reshape(
B, 2 * self.num_heads, self.head_dim, self.local_len, N).permute(0, 1, 4, 2, 3).chunk(2, dim=1)
# Compute local similarity
attn_local = ((q_norm_scaled.unsqueeze(-2) @ k_local).squeeze(-2) \
+ self.relative_pos_bias_local.unsqueeze(1)).masked_fill(self.padding_mask, float('-inf'))
# Generate pooled features
x_ = x.permute(0, 2, 1).reshape(B, -1, H, W).contiguous()
x_ = self.pool(self.act(self.sr(x_))).reshape(B, -1, self.pool_len).permute(0, 2, 1)
x_ = self.norm(x_)
# Generate pooled keys and values
kv_pool = self.kv(x_).reshape(B, self.pool_len, 2 * self.num_heads, self.head_dim).permute(0, 2, 1, 3)
k_pool, v_pool = kv_pool.chunk(2, dim=1)
# Use MLP to generate continuous relative positional bias for pooled features.
pool_bias = self.cpb_fc2(self.cpb_act(self.cpb_fc1(relative_coords_table))).transpose(0, 1)[:,
relative_pos_index.view(-1)].view(-1, N, self.pool_len)
# Compute pooled similarity
attn_pool = q_norm_scaled @ F.normalize(k_pool, dim=-1).transpose(-2, -1) + pool_bias
# Concatenate local & pooled similarity matrices and calculate attention weights through the same Softmax
attn = torch.cat([attn_local, attn_pool], dim=-1).softmax(dim=-1)
attn = self.attn_drop(attn)
# Split the attention weights and separately aggregate the values of local & pooled features
attn_local, attn_pool = torch.split(attn, [self.local_len, self.pool_len], dim=-1)
x_local = (((q_norm @ self.learnable_tokens) + self.learnable_bias + attn_local).unsqueeze(
-2) @ v_local.transpose(-2, -1)).squeeze(-2)
x_pool = attn_pool @ v_pool
x = (x_local + x_pool).transpose(1, 2).reshape(B, N, C)
# Linear projection and output
x = self.proj(x)
x = self.proj_drop(x)
return x
class Attention(nn.Module):
def __init__(self, dim, input_resolution, num_heads=8, qkv_bias=True, attn_drop=0., proj_drop=0.):
super().__init__()
assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.temperature = nn.Parameter(
torch.log((torch.ones(num_heads, 1, 1) / 0.24).exp() - 1)) # Initialize softplus(temperature) to 1/0.24.
# Generate sequnce length scale
self.register_buffer("seq_length_scale", torch.as_tensor(np.log(input_resolution[0] * input_resolution[1])),
persistent=False)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.query_embedding = nn.Parameter(
nn.init.trunc_normal_(torch.empty(self.num_heads, 1, self.head_dim), mean=0, std=0.02))
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
# mlp to generate continuous relative position bias
self.cpb_fc1 = nn.Linear(2, 512, bias=True)
self.cpb_act = nn.ReLU(inplace=True)
self.cpb_fc2 = nn.Linear(512, num_heads, bias=True)
def forward(self, x, H, W, relative_pos_index, relative_coords_table):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, -1, 3 * self.num_heads, self.head_dim).permute(0, 2, 1, 3)
q, k, v = qkv.chunk(3, dim=1)
# Use MLP to generate continuous relative positional bias
rel_bias = self.cpb_fc2(self.cpb_act(self.cpb_fc1(relative_coords_table))).transpose(0, 1)[:,
relative_pos_index.view(-1)].view(-1, N, N)
# Calculate attention map using sequence length scaled cosine attention and query embedding
attn = ((F.normalize(q, dim=-1) + self.query_embedding) * F.softplus(
self.temperature) * self.seq_length_scale) @ F.normalize(k, dim=-1).transpose(-2, -1) + rel_bias
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
class Block(nn.Module):
def __init__(self, dim, num_heads, input_resolution, window_size=3, mlp_ratio=4.,
qkv_bias=False, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
super().__init__()
self.norm1 = norm_layer(dim)
if sr_ratio == 1:
self.attn = Attention(
dim,
input_resolution,
num_heads=num_heads,
qkv_bias=qkv_bias,
attn_drop=attn_drop,
proj_drop=drop)
else:
self.attn = AggregatedAttention(
dim,
input_resolution,
window_size=window_size,
num_heads=num_heads,
qkv_bias=qkv_bias,
attn_drop=attn_drop,
proj_drop=drop,
sr_ratio=sr_ratio)
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = ConvolutionalGLU(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
# 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()
def forward(self, x, H, W, relative_pos_index, relative_coords_table):
x = x + self.drop_path(self.attn(self.norm1(x), H, W, relative_pos_index, relative_coords_table))
x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
return x
class OverlapPatchEmbed(nn.Module):
""" Image to Patch Embedding
"""
def __init__(self, patch_size=7, stride=4, in_chans=3, embed_dim=768):
super().__init__()
patch_size = to_2tuple(patch_size)
assert max(patch_size) > stride, "Set larger patch_size than stride"
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
padding=(patch_size[0] // 2, patch_size[1] // 2))
self.norm = nn.LayerNorm(embed_dim)
def forward(self, x):
x = self.proj(x)
_, _, H, W = x.shape
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
return x, H, W
class TransNeXt(nn.Module):
'''
The parameter "img size" is primarily utilized for generating relative spatial coordinates,
which are used to compute continuous relative positional biases. As this TransNeXt implementation does not support multi-scale inputs,
it is recommended to set the "img size" parameter to a value that is exactly the same as the resolution of the inference images.
It is not advisable to set the "img size" parameter to a value exceeding 800x800.
The "pretrain size" refers to the "img size" used during the initial pre-training phase,
which is used to scale the relative spatial coordinates for better extrapolation by the MLP.
For models trained on ImageNet-1K at a resolution of 224x224,
as well as downstream task models fine-tuned based on these pre-trained weights,
the "pretrain size" parameter should be set to 224x224.
'''
def __init__(self, img_size=640, pretrain_size=None, window_size=[3, 3, 3, None],
patch_size=16, in_chans=3, num_classes=1000, embed_dims=[64, 128, 256, 512],
num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, drop_rate=0.,
attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], num_stages=4):
super().__init__()
self.num_classes = num_classes
self.depths = depths
self.num_stages = num_stages
pretrain_size = pretrain_size or img_size
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
cur = 0
for i in range(num_stages):
# Generate relative positional coordinate table and index for each stage to compute continuous relative positional bias.
relative_pos_index, relative_coords_table = get_relative_position_cpb(
query_size=to_2tuple(img_size // (2 ** (i + 2))),
key_size=to_2tuple(img_size // (2 ** (num_stages + 1))),
pretrain_size=to_2tuple(pretrain_size // (2 ** (i + 2))))
self.register_buffer(f"relative_pos_index{i + 1}", relative_pos_index, persistent=False)
self.register_buffer(f"relative_coords_table{i + 1}", relative_coords_table, persistent=False)
patch_embed = OverlapPatchEmbed(patch_size=patch_size * 2 - 1 if i == 0 else 3,
stride=patch_size if i == 0 else 2,
in_chans=in_chans if i == 0 else embed_dims[i - 1],
embed_dim=embed_dims[i])
block = nn.ModuleList([Block(
dim=embed_dims[i], input_resolution=to_2tuple(img_size // (2 ** (i + 2))), window_size=window_size[i],
num_heads=num_heads[i], mlp_ratio=mlp_ratios[i], qkv_bias=qkv_bias,
drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + j], norm_layer=norm_layer,
sr_ratio=sr_ratios[i])
for j in range(depths[i])])
norm = norm_layer(embed_dims[i])
cur += depths[i]
setattr(self, f"patch_embed{i + 1}", patch_embed)
setattr(self, f"block{i + 1}", block)
setattr(self, f"norm{i + 1}", norm)
for n, m in self.named_modules():
self._init_weights(m, n)
self.width_list = [i.size(1) for i in self.forward(torch.randn(1, 3, 640, 640))]
def _init_weights(self, m: nn.Module, name: str = ''):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Conv2d):
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
fan_out //= m.groups
m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
nn.init.zeros_(m.bias)
nn.init.ones_(m.weight)
def forward(self, x):
B = x.shape[0]
feature = []
for i in range(self.num_stages):
patch_embed = getattr(self, f"patch_embed{i + 1}")
block = getattr(self, f"block{i + 1}")
norm = getattr(self, f"norm{i + 1}")
x, H, W = patch_embed(x)
relative_pos_index = getattr(self, f"relative_pos_index{i + 1}")
relative_coords_table = getattr(self, f"relative_coords_table{i + 1}")
for blk in block:
x = blk(x, H, W, relative_pos_index.to(x.device), relative_coords_table.to(x.device))
x = norm(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
feature.append(x)
return feature
class transnext_micro(TransNeXt):
def __init__(self, **kwargs):
super().__init__(window_size=[3, 3, 3, None],
patch_size=4, embed_dims=[48, 96, 192, 384], num_heads=[2, 4, 8, 16],
mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 15, 2], sr_ratios=[8, 4, 2, 1])
class transnext_tiny(TransNeXt):
def __init__(self, **kwargs):
super().__init__(window_size=[3, 3, 3, None],
patch_size=4, embed_dims=[72, 144, 288, 576], num_heads=[3, 6, 12, 24],
mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 15, 2], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.3)
class transnext_small(TransNeXt):
def __init__(self, **kwargs):
super().__init__(window_size=[3, 3, 3, None],
patch_size=4, embed_dims=[72, 144, 288, 576], num_heads=[3, 6, 12, 24],
mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[5, 5, 22, 5], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.5)
class transnext_base(TransNeXt):
def __init__(self, **kwargs):
super().__init__(window_size=[3, 3, 3, None],
patch_size=4, embed_dims=[96, 192, 384, 768], num_heads=[4, 8, 16, 32],
mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[5, 5, 23, 5], sr_ratios=[8, 4, 2, 1],
drop_rate=0.0, drop_path_rate=0.6,)
下面教大家如何修改该网络结构,主干网络结构的修改步骤比较复杂,我也会将task.py文件上传到CSDN的文件中,大家如果自己修改不正确,可以尝试用我的task.py文件替换你的,然后只需要修改其中的第1、2、3、5步即可。
?修改过程中大家一定要仔细?
首先我门中到如下“ultralytics/nn”的目录,我们在这个目录下在创建一个新的目录,名字为'Addmodules'(此文件之后就用于存放我们的所有改进机制),之后我们在创建的目录内创建一个新的py文件复制粘贴进去 ,可以根据文章改进机制来起,这里大家根据自己的习惯命名即可。
第二步我们在我们创建的目录内创建一个新的py文件名字为'__init__.py'(只需要创建一个即可),然后在其内部导入我们本文的改进机制即可,其余代码均为未发大家没有不用理会!。
第三步我门中到如下文件'ultralytics/nn/tasks.py'然后在开头导入我们的所有改进机制(如果你用了我多个改进机制,这一步只需要修改一次即可)。
添加如下两行代码!!!
?
找到七百多行大概把具体看图片,按照图片来修改就行,添加红框内的部分,注意没有()只是函数名(此处我的文件里已经添加很多了后期都会发出来,大家没有的不用理会即可)。
elif m in {自行添加对应的模型即可,下面都是一样的}:
m = m(*args)
c2 = m.width_list # 返回通道列表
backbone = True
用下面的代码替换红框内的内容。?
?
if isinstance(c2, list):
m_ = m
m_.backbone = True
else:
m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
t = str(m)[8:-2].replace('__main__.', '') # module type
m.np = sum(x.numel() for x in m_.parameters()) # number params
m_.i, m_.f, m_.type = i + 4 if backbone else i, f, t # attach index, 'from' index, type
if verbose:
LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f} {t:<45}{str(args):<30}') # print
save.extend(
x % (i + 4 if backbone else i) for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
layers.append(m_)
if i == 0:
ch = []
if isinstance(c2, list):
ch.extend(c2)
if len(c2) != 5:
ch.insert(0, 0)
else:
ch.append(c2)
修改七这里非常要注意,不是文件开头YOLOv8的那predict,是400+行的RTDETR的predict!!!初始模型如下,用我给的代码替换即可!!!
?
代码如下->
def predict(self, x, profile=False, visualize=False, batch=None, augment=False, embed=None):
"""
Perform a forward pass through the model.
Args:
x (torch.Tensor): The input tensor.
profile (bool, optional): If True, profile the computation time for each layer. Defaults to False.
visualize (bool, optional): If True, save feature maps for visualization. Defaults to False.
batch (dict, optional): Ground truth data for evaluation. Defaults to None.
augment (bool, optional): If True, perform data augmentation during inference. Defaults to False.
embed (list, optional): A list of feature vectors/embeddings to return.
Returns:
(torch.Tensor): Model's output tensor.
"""
y, dt, embeddings = [], [], [] # outputs
for m in self.model[:-1]: # except the head part
if m.f != -1: # if not from previous layer
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
if profile:
self._profile_one_layer(m, x, dt)
if hasattr(m, 'backbone'):
x = m(x)
if len(x) != 5: # 0 - 5
x.insert(0, None)
for index, i in enumerate(x):
if index in self.save:
y.append(i)
else:
y.append(None)
x = x[-1] # 最后一个输出传给下一层
else:
x = m(x) # run
y.append(x if m.i in self.save else None) # save output
if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
if embed and m.i in embed:
embeddings.append(nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1)) # flatten
if m.i == max(embed):
return torch.unbind(torch.cat(embeddings, 1), dim=0)
head = self.model[-1]
x = head([y[j] for j in head.f], batch) # head inference
return x
我们将下面的s用640替换即可,这一步也是部分的主干可以不修改,但有的不修改就会报错,所以我们还是修改一下。
计算的GFLOPs计算异常不打印,所以需要额外修改一处,?我们找到如下文件'ultralytics/utils/torch_utils.py'文件内有如下的代码按照如下的图片进行修改,大家看好函数就行,其中红框的640可能和你的不一样, 然后用我给的代码替换掉整个代码即可。
def get_flops(model, imgsz=640):
"""Return a YOLO model's FLOPs."""
try:
model = de_parallel(model)
p = next(model.parameters())
# stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32 # max stride
stride = 640
im = torch.empty((1, 3, stride, stride), device=p.device) # input image in BCHW format
flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1E9 * 2 if thop else 0 # stride GFLOPs
imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz] # expand if int/float
return flops * imgsz[0] / stride * imgsz[1] / stride # 640x640 GFLOPs
except Exception:
return 0
有些读者的数据集部分图片比较特殊,在验证的时候会导致形状不匹配的报错,如果大家在验证的时候报错形状不匹配的错误可以固定验证集的图片尺寸,方法如下?->
找到下面这个文件ultralytics/models/yolo/detect/train.py然后其中有一个类是DetectionTrainer class中的build_dataset函数中的一个参数rect=mode == 'val'改为rect=False
大家复制下面的yaml文件,然后通过我给大家的运行代码运行即可,RT-DETR的调参部分需要后面的文章给大家讲,现在目前免费给大家看这一部分不开放。
# 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, transnext_micro, []] # 4
head:
- [-1, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 5 input_proj.2
- [-1, 1, AIFI, [1024, 8]] # 6
- [-1, 1, Conv, [256, 1, 1]] # 7, Y5, lateral_convs.0
- [-1, 1, nn.Upsample, [None, 2, 'nearest']] # 8
- [3, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 9 input_proj.1
- [[-2, -1], 1, Concat, [1]] # 10
- [-1, 3, RepC3, [256, 0.5]] # 11, fpn_blocks.0
- [-1, 1, Conv, [256, 1, 1]] # 12, Y4, lateral_convs.1
- [-1, 1, nn.Upsample, [None, 2, 'nearest']] # 13
- [2, 1, Conv, [256, 1, 1, None, 1, 1, False]] # 14 input_proj.0
- [[-2, -1], 1, Concat, [1]] # 15 cat backbone P4
- [-1, 3, RepC3, [256, 0.5]] # X3 (16), fpn_blocks.1
- [-1, 1, Conv, [256, 3, 2]] # 17, downsample_convs.0
- [[-1, 12], 1, Concat, [1]] # 18 cat Y4
- [-1, 3, RepC3, [256, 0.5]] # F4 (19), pan_blocks.0
- [-1, 1, Conv, [256, 3, 2]] # 20, downsample_convs.1
- [[-1, 7], 1, Concat, [1]] # 21 cat Y5
- [-1, 3, RepC3, [256, 0.5]] # F5 (22), pan_blocks.1
- [[16, 19, 22], 1, RTDETRDecoder, [nc, 256, 300, 4, 8, 3]] # Detect(P3, P4, P5)
大家可以创建一个train.py文件将下面的代码粘贴进去然后替换你的文件运行即可开始训练。
import warnings
from ultralytics import RTDETR
warnings.filterwarnings('ignore')
if __name__ == '__main__':
model = RTDETR('替换你想要运行的yaml文件')
# model.load('') # 可以加载你的版本预训练权重
model.train(data=r'替换你的数据集地址即可',
cache=False,
imgsz=640,
epochs=72,
batch=4,
workers=0,
device='0',
project='runs/RT-DETR-train',
name='exp',
# amp=True
)
下面是成功运行的截图(确保我的改进机制是可用的),已经完成了有1个epochs的训练,图片太大截不全第2个epochs了。?
?
从今天开始正式开始更新RT-DETR剑指论文专栏,本专栏的内容会迅速铺开,在短期呢大量更新,价格也会乘阶梯性上涨,所以想要和我一起学习RT-DETR改进,可以在前期直接关注,本文专栏旨在打造全网最好的RT-DETR专栏为想要发论文的家进行服务。