前言

大家好，我是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在各种视觉任务中，包括图像分类、目标检测和语义分割，都显示出优异的性能。

 专栏链接：RT-DETR剑指论文专栏，持续复现各种顶会内容——论文收割机RT-DETR

目录

 一、本文介绍

二、TransNeXt的框架原理

2.1 聚合注意力机制 

2.2 卷积GLU

2.3  TransNeXt的架构示意图

三、TransNeXt的核心代码 

 四、手把手教你添加TransNeXt机制

4.1 修改一

4.2 修改二 

4.3 修改三 

4.4 修改四

4.5 修改五

4.6 修改六

4.7 修改七 

4.8 修改八

4.9 RT-DETR不能打印计算量问题的解决

4.10 可选修改

五、TransNeXt的yaml文件

5.1 yaml文件

5.2 运行文件

5.3 成功训练截图

六、全文总结

---

二、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张224²分辨率的ImageNet-1K验证集图像，展示了不同方法与流行背景网络之间的有效感受野（ERF）的可视化对比。图中展示了四种不同的注意力机制：局部注意力、滑动窗口注意力与卷积、池化注意力，以及本文提出的像素聚焦注意力，以及它们与生物视觉系统的对比。每种机制下，红星表示当前查询的位置，黑色区域表示当前查询无法感知的区域。右侧的图表比较了TransNeXt（本文提出的方法）和其他几种流行的模型在处理视觉信息时的差异。

---

2.1 聚合注意力机制 

聚合注意力机制（Aggregated Attention, AA）是TransNeXt模型中的一个核心创新，它融合了多种注意力机制并为多尺度输入增强了外推能力，具体包括以下几点：

1. 像素聚焦注意力：

• 该机制受到生物视觉系统功能的启发，旨在为每个查询（query）提供细粒度的感知，同时保持对全局信息的粗粒度认识。
• 通过使用双路径设计，结合了以查询为中心的滑动窗口注意力和池化注意力，实现了像素级的平移等变性，模拟眼球运动的特性。
• 这种设计导致细粒度和粗粒度特征之间的竞争，使得像素聚焦注意力转变为一个多尺度的注意力机制。

2. 集成多样的注意力机制：

• 研究表明，将可学习的查询前缀整合到注意力机制中，并直接对其进行优化，对于图像分类、目标检测和语义分割等明确定义的任务是有效且高效的。
• 添加可学习的查询嵌入到所有传统的查询-键值-值（QKV）注意力中的查询令牌可以实现类似的信息聚合效果，且额外开销微不足道。

3. 位置注意力：

• 使用一组可学习的键与输入中的查询互动，获得注意力权重，即查询-可学习值（QLV）注意力。
• 与传统的QKV注意力相比，该方法打破了键和值之间一对一的对应关系，可以学习当前查询的更隐式的相对位置信息。

4. 多尺度输入的外推能力：

• 为了克服多尺度输入的问题，提出了长度缩放的余弦注意力，该方法使用余弦相似性，通过增加一个额外的可学习系数来改善训练大型视觉模型的稳定性。
• 这种设计有助于处理随着输入序列长度增加而减少的注意力输出的置信度问题，并通过长度缩放来维持熵的不变性，以便更好地推广到未知长度。

5. 聚合注意力（Aggregated Attention）：

• 通过应用上述多样的注意力聚合方法和技术，提出了增强版的像素聚焦注意力，即聚合像素聚焦注意力。
• 该机制不仅聚合了多种注意力机制，还在计算过程中引入了长度缩放的余弦注意力和可学习的查询嵌入，以及特定的位置编码方法，进一步提升了对多尺度输入的处理能力。

总结：聚合注意力机制通过模拟生物视觉系统，提供了一种更自然的视觉感知方式，可以有效地处理来自不同层次和尺度的信息，并通过结合不同的注意力路径和可学习组件，增强了模型对于多尺度输入的外推能力。

缺点：但是它总结了这么多注意力机制，它限制了通道数减少了参数量，但是其运算非常复杂导致速度很慢

图3是像素聚焦注意力（左图）与聚合注意力（右图）之间对比的插图。两者都具有10x10的特征尺寸，一个3x3的窗口尺寸和2x2的池化尺寸。

左图（像素聚焦注意力）：

• 展示了一个包含滑动窗口注意力来收集局部信息的结构。这涉及到对最近令牌的查询和键值对比较以及位置偏差的应用。
• 使用池化操作来收集更广泛区域的信息，通过AxialPool来实现，这样的设计旨在捕捉到全局信息。
• 在计算完注意力权重之后，这些信息会被合并并通过AxialPool、LayerNorm等操作处理，最终生成输出特征图。

右图（聚合注意力）：

• 在像素聚焦注意力的基础上增加了一些关键的组件来形成聚合注意力。引入了位置注意力，它作为动态相对位置偏差使用，和可学习的令牌相结合，来增强模型对位置的感知能力。
• 加入了查询嵌入（Query Embedding），这是一种改进，它使得每个查询都与一个额外的可学习向量相结合，以进一步优化注意力权重的计算。
• 同样地，通过各种层操作处理后生成输出特征图。

---

2.2 卷积GLU

卷积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）：

• 输入通过一个线性层，接着是一个深度卷积层（DW Conv 3x3），然后是一个激活函数，再是一个线性层。
• 最后，输入和线性层的输出相加，形成最终的输出。

3. 门控线性单元（Gated Linear Unit, GLU）：

• 输入通过两个平行的线性层，一个直接输出，另一个先经过激活函数，然后输出。
• 这两个输出进行逐元素乘法操作，然后通过另一个线性层。
• 最后，输入和这个线性层的输出相加，形成最终的输出。

4. 带有SE模块的前馈网络（FFN with SE module）：

• 输入通过一个线性层，接着是激活函数，然后是另一个线性层。
• 同时，输入经过全局平均池化，然后是一个线性层，ReLU激活函数，另一个线性层，以及Sigmoid函数，形成SE模块的输出。
• SE模块的输出与前馈网络的中间输出进行逐元素乘法操作。
• 最后，输入和乘法操作后的输出相加，形成最终的输出。

5. 卷积门控线性单元（Convolutional Gated Linear Unit）：

• 输入通过一个线性层，接着是一个深度卷积层（DW Conv 3x3），然后是激活函数。
• 同时，输入也经过另外一个线性层的处理。
• 这两个部分的输出进行逐元素乘法操作。
• 最后，输入和乘法操作后的输出相加，形成最终的输出。

---

2.3  TransNeXt的架构示意图

图5展示了TransNeXt架构的一个示意图，揭示了其内部的组件和数据流。这个架构通过多个阶段的处理来处理输入图像，其中每个阶段都包含卷积GLU和聚合注意力机制的层。以下是每个阶段的详细介绍：

1. 图像输入：

• 输入图像的维度是，其中H和W是图像的高度和宽度，3代表RGB三个颜色通道。

2. 阶段1：

• 首先，图像通过一个补丁嵌入层（Patch Embedding），这个层将图像分割成更小的块，并将每块映射成一个向量，向量的大小由  决定，这里的C是嵌入向量的维度。
• 接着，数据流通过多个卷积GLU和聚合注意力机制的层，每个层后都跟随一个层归一化。
• 这个阶段重复 次，每次都可能对特征图进行下采样，减少其空间维度并增加通道数例如。

3. 阶段2和3：

• 这些阶段与阶段1类似，但是每个阶段都会进一步减少特征图的空间维度并增加通道数（例如阶段2是，阶段3是。
• 在这些阶段中，模型继续使用卷积GLU和聚合注意力来处理和提炼特征，这些特征对应于更抽象的图像表示。
• 阶段2和3分别重复   和次。

4. 阶段4：

• 在最后一个阶段，模型增加了多头自注意力（Multi-Head Self-Attention）层，这是标准Transformer架构的关键部分，它可以捕捉不同头部间的不同表示。
• 同样，这个阶段还使用卷积GLU和层归一化，重复 次。

总结：通过这些阶段的处理，TransNeXt模型能够逐步提取和处理图像特征，从局部像素级特征到更高层次的抽象表示。每个阶段的输出都准备好进入下一个阶段，直到最终生成能够用于图像分类、目标检测或语义分割任务的高级特征。此架构展示了如何通过结合卷积和注意力机制来有效地处理视觉数据，同时逐步增加通道数和降低空间分辨率，以提高计算效率和模型性能。

---

三、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,)

四、手把手教你添加TransNeXt机制

下面教大家如何修改该网络结构，主干网络结构的修改步骤比较复杂，我也会将task.py文件上传到CSDN的文件中，大家如果自己修改不正确，可以尝试用我的task.py文件替换你的，然后只需要修改其中的第1、2、3、5步即可。

⭐修改过程中大家一定要仔细⭐

---

4.1 修改一

首先我门中到如下“ultralytics/nn”的目录，我们在这个目录下在创建一个新的目录，名字为'Addmodules'（此文件之后就用于存放我们的所有改进机制），之后我们在创建的目录内创建一个新的py文件复制粘贴进去 ，可以根据文章改进机制来起，这里大家根据自己的习惯命名即可。

---

4.2 修改二 

第二步我们在我们创建的目录内创建一个新的py文件名字为'__init__.py'（只需要创建一个即可），然后在其内部导入我们本文的改进机制即可，其余代码均为未发大家没有不用理会！。

---

4.3 修改三 

第三步我门中到如下文件'ultralytics/nn/tasks.py'然后在开头导入我们的所有改进机制（如果你用了我多个改进机制，这一步只需要修改一次即可）。

---

4.4 修改四

添加如下两行代码！！！

​

---

4.5 修改五

找到七百多行大概把具体看图片，按照图片来修改就行，添加红框内的部分，注意没有()只是函数名（此处我的文件里已经添加很多了后期都会发出来，大家没有的不用理会即可）。

        elif m in {自行添加对应的模型即可，下面都是一样的}:
            m = m(*args)
            c2 = m.width_list  # 返回通道列表
            backbone = True

---

4.6 修改六

用下面的代码替换红框内的内容。 

​

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)

---

4.7 修改七 

修改七这里非常要注意，不是文件开头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

4.8 修改八

我们将下面的s用640替换即可，这一步也是部分的主干可以不修改，但有的不修改就会报错，所以我们还是修改一下。

---

4.9 RT-DETR不能打印计算量问题的解决

计算的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

---

4.10 可选修改

有些读者的数据集部分图片比较特殊，在验证的时候会导致形状不匹配的报错，如果大家在验证的时候报错形状不匹配的错误可以固定验证集的图片尺寸，方法如下 ->

找到下面这个文件ultralytics/models/yolo/detect/train.py然后其中有一个类是DetectionTrainer class中的build_dataset函数中的一个参数rect=mode == 'val'改为rect=False

---

五、TransNeXt的yaml文件

5.1 yaml文件

大家复制下面的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)

---

5.2 运行文件

大家可以创建一个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
                )

---

5.3 成功训练截图

下面是成功运行的截图（确保我的改进机制是可用的），已经完成了有1个epochs的训练，图片太大截不全第2个epochs了。 

​

---

六、全文总结

从今天开始正式开始更新RT-DETR剑指论文专栏，本专栏的内容会迅速铺开，在短期呢大量更新，价格也会乘阶梯性上涨，所以想要和我一起学习RT-DETR改进，可以在前期直接关注，本文专栏旨在打造全网最好的RT-DETR专栏为想要发论文的家进行服务。

 专栏链接：RT-DETR剑指论文专栏，持续复现各种顶会内容——论文收割机RT-DETR