学习资源站

YOLOv11改进-主干_Backboen篇-目标检测网络最新成果TransNeXt特征提取网络（yolov11全系列可用轻量化）

一、本文介绍

本文给大家带来的改进机制是 TransNeXt特征提取网络 ,将其应用在我们的 特征提取网络 来提取特征，同时本文给大家解决其自带的一个报错，通过结合聚合的像素聚焦注意力和卷积GLU，模拟生物视觉系统，特别是对于中心凹的视觉感知。这种方法使得每个像素都能实现全局感知，并强化了 模型 的信息混合和自然视觉感知能力。TransNeXt在 各种视觉任务中，包括图像分类、目标检测和语义分割 ，都显示出优异的性能(该模型的训练时间很长这是需要大家注意的)。

欢迎大家订阅我的专栏一起学习YOLO！

（本文内容可根据yolov11的N、S、M、L、X进行二次缩放，轻量化更上一层）。

二、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 \times W \times 3 ，其中H和W是图像的高度和宽度，3代表RGB三个颜色通道。

2. 阶段1：

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

3. 阶段2和3：

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

4. 阶段4：

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

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

三、TransNeXt的核心代码

官方的代码提供了两种一种是需要编译的，条件非常苛刻，但是速度运行比不编译的要快一些，我这里提供的是不编译的方便大家运行，大家要是想用编译的可以从我给的代码链接中找到，按照我的对比着来进行修改即可。 

  1. import torch
  2. import torch.nn as nn
  3. import torch.nn.functional as F
  4. import numpy as np
  5. from functools import partial
  6. from timm.models.layers import DropPath, to_2tuple, trunc_normal_
  7. import math
  9. __all__ = ['transnext_micro', 'transnext_small', 'transnext_base', 'transnext_tiny']
  11. class DWConv(nn.Module):
  12.     def __init__(self, dim=768):
  13.         super(DWConv, self).__init__()
  14.         self.dwconv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=True, groups=dim)
  16.     def forward(self, x, H, W):
  17.         B, N, C = x.shape
  18.         x = x.transpose(1, 2).view(B, C, H, W).contiguous()
  19.         x = self.dwconv(x)
  20.         x = x.flatten(2).transpose(1, 2)
  22.         return x
  25. class ConvolutionalGLU(nn.Module):
  26.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
  27.         super().__init__()
  28.         out_features = out_features or in_features
  29.         hidden_features = hidden_features or in_features
  30.         hidden_features = int(2 * hidden_features / 3)
  31.         self.fc1 = nn.Linear(in_features, hidden_features * 2)
  32.         self.dwconv = DWConv(hidden_features)
  33.         self.act = act_layer()
  34.         self.fc2 = nn.Linear(hidden_features, out_features)
  35.         self.drop = nn.Dropout(drop)
  37.     def forward(self, x, H, W):
  38.         x, v = self.fc1(x).chunk(2, dim=-1)
  39.         x = self.act(self.dwconv(x, H, W)) * v
  40.         x = self.drop(x)
  41.         x = self.fc2(x)
  42.         x = self.drop(x)
  43.         return x
  46. @torch.no_grad()
  47. def get_relative_position_cpb(query_size, key_size, pretrain_size=None):
  48.     # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  49.     pretrain_size = pretrain_size or query_size
  50.     axis_qh = torch.arange(query_size[0], dtype=torch.float32)
  51.     axis_kh = F.adaptive_avg_pool1d(axis_qh.unsqueeze(0), key_size[0]).squeeze(0)
  52.     axis_qw = torch.arange(query_size[1], dtype=torch.float32)
  53.     axis_kw = F.adaptive_avg_pool1d(axis_qw.unsqueeze(0), key_size[1]).squeeze(0)
  54.     axis_kh, axis_kw = torch.meshgrid(axis_kh, axis_kw)
  55.     axis_qh, axis_qw = torch.meshgrid(axis_qh, axis_qw)
  57.     axis_kh = torch.reshape(axis_kh, [-1])
  58.     axis_kw = torch.reshape(axis_kw, [-1])
  59.     axis_qh = torch.reshape(axis_qh, [-1])
  60.     axis_qw = torch.reshape(axis_qw, [-1])
  62.     relative_h = (axis_qh[:, None] - axis_kh[None, :]) / (pretrain_size[0] - 1) * 8
  63.     relative_w = (axis_qw[:, None] - axis_kw[None, :]) / (pretrain_size[1] - 1) * 8
  64.     relative_hw = torch.stack([relative_h, relative_w], dim=-1).view(-1, 2)
  66.     relative_coords_table, idx_map = torch.unique(relative_hw, return_inverse=True, dim=0)
  68.     relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
  69.         torch.abs(relative_coords_table) + 1.0) / torch.log2(torch.tensor(8, dtype=torch.float32))
  71.     return idx_map, relative_coords_table
  74. @torch.no_grad()
  75. def get_seqlen_and_mask(input_resolution, window_size):
  76.     attn_map = F.unfold(torch.ones([1, 1, input_resolution[0], input_resolution[1]]), window_size,
  77.                         dilation=1, padding=(window_size // 2, window_size // 2), stride=1)
  78.     attn_local_length = attn_map.sum(-2).squeeze().unsqueeze(-1)
  79.     attn_mask = (attn_map.squeeze(0).permute(1, 0)) == 0
  80.     return attn_local_length, attn_mask
  83. class AggregatedAttention(nn.Module):
  84.     def __init__(self, dim, input_resolution, num_heads=8, window_size=3, qkv_bias=True,
  85.                  attn_drop=0., proj_drop=0., sr_ratio=1):
  86.         super().__init__()
  87.         assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
  89.         self.dim = dim
  90.         self.num_heads = num_heads
  91.         self.head_dim = dim // num_heads
  93.         self.sr_ratio = sr_ratio
  95.         assert window_size % 2 == 1, "window size must be odd"
  96.         self.window_size = window_size
  97.         self.local_len = window_size ** 2
  99.         self.pool_H, self.pool_W = input_resolution[0] // self.sr_ratio, input_resolution[1] // self.sr_ratio
  100.         self.pool_len = self.pool_H * self.pool_W
  102.         self.unfold = nn.Unfold(kernel_size=window_size, padding=window_size // 2, stride=1)
  103.         self.temperature = nn.Parameter(
  104.             torch.log((torch.ones(num_heads, 1, 1) / 0.24).exp() - 1))  # Initialize softplus(temperature) to 1/0.24.
  106.         self.q = nn.Linear(dim, dim, bias=qkv_bias)
  107.         self.query_embedding = nn.Parameter(
  108.             nn.init.trunc_normal_(torch.empty(self.num_heads, 1, self.head_dim), mean=0, std=0.02))
  109.         self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
  110.         self.attn_drop = nn.Dropout(attn_drop)
  111.         self.proj = nn.Linear(dim, dim)
  112.         self.proj_drop = nn.Dropout(proj_drop)
  114.         # Components to generate pooled features.
  115.         self.pool = nn.AdaptiveAvgPool2d((self.pool_H, self.pool_W))
  116.         self.sr = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0)
  117.         self.norm = nn.LayerNorm(dim)
  118.         self.act = nn.GELU()
  120.         # mlp to generate continuous relative position bias
  121.         self.cpb_fc1 = nn.Linear(2, 512, bias=True)
  122.         self.cpb_act = nn.ReLU(inplace=True)
  123.         self.cpb_fc2 = nn.Linear(512, num_heads, bias=True)
  125.         # relative bias for local features
  126.         self.relative_pos_bias_local = nn.Parameter(
  127.             nn.init.trunc_normal_(torch.empty(num_heads, self.local_len), mean=0,
  128.                                   std=0.0004))
  130.         # Generate padding_mask && sequnce length scale
  131.         local_seq_length, padding_mask = get_seqlen_and_mask(input_resolution, window_size)
  132.         self.register_buffer("seq_length_scale", torch.as_tensor(np.log(local_seq_length.numpy() + self.pool_len)),
  133.                              persistent=False)
  134.         self.register_buffer("padding_mask", padding_mask, persistent=False)
  136.         # dynamic_local_bias:
  137.         self.learnable_tokens = nn.Parameter(
  138.             nn.init.trunc_normal_(torch.empty(num_heads, self.head_dim, self.local_len), mean=0, std=0.02))
  139.         self.learnable_bias = nn.Parameter(torch.zeros(num_heads, 1, self.local_len))
  141.     def forward(self, x, H, W, relative_pos_index, relative_coords_table):
  142.         B, N, C = x.shape
  144.         # Generate queries, normalize them with L2, add query embedding, and then magnify with sequence length scale and temperature.
  145.         # Use softplus function ensuring that the temperature is not lower than 0.
  146.         q_norm = F.normalize(self.q(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3), dim=-1)
  147.         q_norm_scaled = (q_norm + self.query_embedding) * F.softplus(self.temperature) * self.seq_length_scale
  149.         # Generate unfolded keys and values and l2-normalize them
  150.         k_local, v_local = self.kv(x).chunk(2, dim=-1)
  151.         k_local = F.normalize(k_local.reshape(B, N, self.num_heads, self.head_dim), dim=-1).reshape(B, N, -1)
  152.         kv_local = torch.cat([k_local, v_local], dim=-1).permute(0, 2, 1).reshape(B, -1, H, W)
  153.         k_local, v_local = self.unfold(kv_local).reshape(
  154.             B, 2 * self.num_heads, self.head_dim, self.local_len, N).permute(0, 1, 4, 2, 3).chunk(2, dim=1)
  156.         # Compute local similarity
  157.         attn_local = ((q_norm_scaled.unsqueeze(-2) @ k_local).squeeze(-2) \
  158.                       + self.relative_pos_bias_local.unsqueeze(1)).masked_fill(self.padding_mask, float('-inf'))
  160.         # Generate pooled features
  161.         x_ = x.permute(0, 2, 1).reshape(B, -1, H, W).contiguous()
  162.         x_ = self.pool(self.act(self.sr(x_))).reshape(B, -1, self.pool_len).permute(0, 2, 1)
  163.         x_ = self.norm(x_)
  165.         # Generate pooled keys and values
  166.         kv_pool = self.kv(x_).reshape(B, self.pool_len, 2 * self.num_heads, self.head_dim).permute(0, 2, 1, 3)
  167.         k_pool, v_pool = kv_pool.chunk(2, dim=1)
  169.         # Use MLP to generate continuous relative positional bias for pooled features.
  170.         pool_bias = self.cpb_fc2(self.cpb_act(self.cpb_fc1(relative_coords_table))).transpose(0, 1)[:,
  171.                     relative_pos_index.view(-1)].view(-1, N, self.pool_len)
  172.         # Compute pooled similarity
  173.         attn_pool = q_norm_scaled @ F.normalize(k_pool, dim=-1).transpose(-2, -1) + pool_bias
  175.         # Concatenate local & pooled similarity matrices and calculate attention weights through the same Softmax
  176.         attn = torch.cat([attn_local, attn_pool], dim=-1).softmax(dim=-1)
  177.         attn = self.attn_drop(attn)
  179.         # Split the attention weights and separately aggregate the values of local & pooled features
  180.         attn_local, attn_pool = torch.split(attn, [self.local_len, self.pool_len], dim=-1)
  181.         x_local = (((q_norm @ self.learnable_tokens) + self.learnable_bias + attn_local).unsqueeze(
  182.             -2) @ v_local.transpose(-2, -1)).squeeze(-2)
  183.         x_pool = attn_pool @ v_pool
  184.         x = (x_local + x_pool).transpose(1, 2).reshape(B, N, C)
  186.         # Linear projection and output
  187.         x = self.proj(x)
  188.         x = self.proj_drop(x)
  190.         return x
  193. class Attention(nn.Module):
  194.     def __init__(self, dim, input_resolution, num_heads=8, qkv_bias=True, attn_drop=0., proj_drop=0.):
  195.         super().__init__()
  196.         assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
  198.         self.dim = dim
  199.         self.num_heads = num_heads
  200.         self.head_dim = dim // num_heads
  201.         self.temperature = nn.Parameter(
  202.             torch.log((torch.ones(num_heads, 1, 1) / 0.24).exp() - 1))  # Initialize softplus(temperature) to 1/0.24.
  203.         # Generate sequnce length scale
  204.         self.register_buffer("seq_length_scale", torch.as_tensor(np.log(input_resolution[0] * input_resolution[1])),
  205.                              persistent=False)
  207.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
  208.         self.query_embedding = nn.Parameter(
  209.             nn.init.trunc_normal_(torch.empty(self.num_heads, 1, self.head_dim), mean=0, std=0.02))
  211.         self.attn_drop = nn.Dropout(attn_drop)
  212.         self.proj = nn.Linear(dim, dim)
  213.         self.proj_drop = nn.Dropout(proj_drop)
  215.         # mlp to generate continuous relative position bias
  216.         self.cpb_fc1 = nn.Linear(2, 512, bias=True)
  217.         self.cpb_act = nn.ReLU(inplace=True)
  218.         self.cpb_fc2 = nn.Linear(512, num_heads, bias=True)
  220.     def forward(self, x, H, W, relative_pos_index, relative_coords_table):
  221.         B, N, C = x.shape
  222.         qkv = self.qkv(x).reshape(B, -1, 3 * self.num_heads, self.head_dim).permute(0, 2, 1, 3)
  223.         q, k, v = qkv.chunk(3, dim=1)
  225.         # Use MLP to generate continuous relative positional bias
  226.         rel_bias = self.cpb_fc2(self.cpb_act(self.cpb_fc1(relative_coords_table))).transpose(0, 1)[:,
  227.                    relative_pos_index.view(-1)].view(-1, N, N)
  229.         # Calculate attention map using sequence length scaled cosine attention and query embedding
  230.         attn = ((F.normalize(q, dim=-1) + self.query_embedding) * F.softplus(
  231.             self.temperature) * self.seq_length_scale) @ F.normalize(k, dim=-1).transpose(-2, -1) + rel_bias
  232.         attn = attn.softmax(dim=-1)
  233.         attn = self.attn_drop(attn)
  234.         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
  235.         x = self.proj(x)
  236.         x = self.proj_drop(x)
  237.         return x
  240. class Block(nn.Module):
  242.     def __init__(self, dim, num_heads, input_resolution, window_size=3, mlp_ratio=4.,
  243.                  qkv_bias=False, drop=0., attn_drop=0.,
  244.                  drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
  245.         super().__init__()
  246.         self.norm1 = norm_layer(dim)
  247.         if sr_ratio == 1:
  248.             self.attn = Attention(
  249.                 dim,
  250.                 input_resolution,
  251.                 num_heads=num_heads,
  252.                 qkv_bias=qkv_bias,
  253.                 attn_drop=attn_drop,
  254.                 proj_drop=drop)
  255.         else:
  256.             self.attn = AggregatedAttention(
  257.                 dim,
  258.                 input_resolution,
  259.                 window_size=window_size,
  260.                 num_heads=num_heads,
  261.                 qkv_bias=qkv_bias,
  262.                 attn_drop=attn_drop,
  263.                 proj_drop=drop,
  264.                 sr_ratio=sr_ratio)
  265.         self.norm2 = norm_layer(dim)
  266.         mlp_hidden_dim = int(dim * mlp_ratio)
  267.         self.mlp = ConvolutionalGLU(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
  269.         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
  270.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  272.     def forward(self, x, H, W, relative_pos_index, relative_coords_table):
  273.         x = x + self.drop_path(self.attn(self.norm1(x), H, W, relative_pos_index, relative_coords_table))
  274.         x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
  276.         return x
  279. class OverlapPatchEmbed(nn.Module):
  280.     """ Image to Patch Embedding
  281.     """
  283.     def __init__(self, patch_size=7, stride=4, in_chans=3, embed_dim=768):
  284.         super().__init__()
  286.         patch_size = to_2tuple(patch_size)
  288.         assert max(patch_size) > stride, "Set larger patch_size than stride"
  289.         self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride,
  290.                               padding=(patch_size[0] // 2, patch_size[1] // 2))
  291.         self.norm = nn.LayerNorm(embed_dim)
  293.     def forward(self, x):
  294.         x = self.proj(x)
  295.         _, _, H, W = x.shape
  296.         x = x.flatten(2).transpose(1, 2)
  297.         x = self.norm(x)
  299.         return x, H, W
  302. class TransNeXt(nn.Module):
  303.     '''
  304.     The parameter "img size" is primarily utilized for generating relative spatial coordinates,
  305.     which are used to compute continuous relative positional biases. As this TransNeXt implementation does not support multi-scale inputs,
  306.     it is recommended to set the "img size" parameter to a value that is exactly the same as the resolution of the inference images.
  307.     It is not advisable to set the "img size" parameter to a value exceeding 800x800.
  308.     The "pretrain size" refers to the "img size" used during the initial pre-training phase,
  309.     which is used to scale the relative spatial coordinates for better extrapolation by the MLP.
  310.     For models trained on ImageNet-1K at a resolution of 224x224,
  311.     as well as downstream task models fine-tuned based on these pre-trained weights,
  312.     the "pretrain size" parameter should be set to 224x224.
  313.     '''
  315.     def __init__(self, factor=0.5, depth_factor=0.5, img_size=640, pretrain_size=None, window_size=[3, 3, 3, None],
  316.                  patch_size=16, in_chans=3, num_classes=1000, embed_dims=[64, 128, 256, 512],
  317.                  num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, drop_rate=0.,
  318.                  attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
  319.                  depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1], num_stages=4):
  320.         super().__init__()
  321.         embed_dims = tuple([int(dim * factor) for dim in embed_dims])
  322.         depths = [max(round(dim * depth_factor), 1) if dim > 1 else dim for dim in depths]
  323.         self.num_classes = num_classes
  324.         self.depths = depths
  325.         self.num_stages = num_stages
  326.         pretrain_size = pretrain_size or img_size
  328.         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
  329.         cur = 0
  331.         for i in range(num_stages):
  332.             # Generate relative positional coordinate table and index for each stage to compute continuous relative positional bias.
  333.             relative_pos_index, relative_coords_table = get_relative_position_cpb(
  334.                 query_size=to_2tuple(img_size // (2 ** (i + 2))),
  335.                 key_size=to_2tuple(img_size // (2 ** (num_stages + 1))),
  336.                 pretrain_size=to_2tuple(pretrain_size // (2 ** (i + 2))))
  338.             self.register_buffer(f"relative_pos_index{i + 1}", relative_pos_index, persistent=False)
  339.             self.register_buffer(f"relative_coords_table{i + 1}", relative_coords_table, persistent=False)
  341.             patch_embed = OverlapPatchEmbed(patch_size=patch_size * 2 - 1 if i == 0 else 3,
  342.                                             stride=patch_size if i == 0 else 2,
  343.                                             in_chans=in_chans if i == 0 else embed_dims[i - 1],
  344.                                             embed_dim=embed_dims[i])
  346.             block = nn.ModuleList([Block(
  347.                 dim=embed_dims[i], input_resolution=to_2tuple(img_size // (2 ** (i + 2))), window_size=window_size[i],
  348.                 num_heads=num_heads[i], mlp_ratio=mlp_ratios[i], qkv_bias=qkv_bias,
  349.                 drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + j], norm_layer=norm_layer,
  350.                 sr_ratio=sr_ratios[i])
  351.                 for j in range(depths[i])])
  352.             norm = norm_layer(embed_dims[i])
  353.             cur += depths[i]
  355.             setattr(self, f"patch_embed{i + 1}", patch_embed)
  356.             setattr(self, f"block{i + 1}", block)
  357.             setattr(self, f"norm{i + 1}", norm)
  359.         for n, m in self.named_modules():
  360.             self._init_weights(m, n)
  362.         self.width_list = [i.size(1) for i in self.forward(torch.randn(1, 3, 640, 640))]
  364.     def _init_weights(self, m: nn.Module, name: str = ''):
  365.         if isinstance(m, nn.Linear):
  366.             trunc_normal_(m.weight, std=.02)
  367.             if m.bias is not None:
  368.                 nn.init.zeros_(m.bias)
  369.         elif isinstance(m, nn.Conv2d):
  370.             fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
  371.             fan_out //= m.groups
  372.             m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
  373.             if m.bias is not None:
  374.                 m.bias.data.zero_()
  375.         elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
  376.             nn.init.zeros_(m.bias)
  377.             nn.init.ones_(m.weight)
  379.     def forward(self, x):
  380.         B = x.shape[0]
  382.         feature = []
  383.         for i in range(self.num_stages):
  384.             patch_embed = getattr(self, f"patch_embed{i + 1}")
  385.             block = getattr(self, f"block{i + 1}")
  386.             norm = getattr(self, f"norm{i + 1}")
  387.             x, H, W = patch_embed(x)
  388.             relative_pos_index = getattr(self, f"relative_pos_index{i + 1}")
  389.             relative_coords_table = getattr(self, f"relative_coords_table{i + 1}")
  390.             for blk in block:
  391.                 x = blk(x, H, W, relative_pos_index.to(x.device), relative_coords_table.to(x.device))
  392.             x = norm(x)
  393.             x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
  394.             feature.append(x)
  396.         return feature
  398. class transnext_micro(TransNeXt):
  399.     def __init__(self, factor, depth_factor, **kwargs):
  400.         super().__init__(factor, depth_factor, window_size=[3, 3, 3, None],
  401.                       patch_size=4, embed_dims=[48, 96, 192, 384], num_heads=[2, 4, 8, 16],
  402.                       mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
  403.                       norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 15, 2], sr_ratios=[8, 4, 2, 1])
  406. class transnext_tiny(TransNeXt):
  407.     def __init__(self,factor, depth_factor,  **kwargs):
  408.         super().__init__(factor, depth_factor, window_size=[3, 3, 3, None],
  409.                          patch_size=4, embed_dims=[72, 144, 288, 576], num_heads=[3, 6, 12, 24],
  410.                          mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
  411.                          norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 15, 2], sr_ratios=[8, 4, 2, 1],
  412.                          drop_rate=0.0, drop_path_rate=0.3)
  415. class transnext_small(TransNeXt):
  416.     def __init__(self, factor, depth_factor, **kwargs):
  417.         super().__init__(factor, depth_factor, window_size=[3, 3, 3, None],
  418.                          patch_size=4, embed_dims=[72, 144, 288, 576], num_heads=[3, 6, 12, 24],
  419.                          mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
  420.                          norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[5, 5, 22, 5], sr_ratios=[8, 4, 2, 1],
  421.                          drop_rate=0.0, drop_path_rate=0.5)
  423. class transnext_base(TransNeXt):
  424.     def __init__(self,factor, depth_factor,  **kwargs):
  425.         super().__init__(factor, depth_factor, window_size=[3, 3, 3, None],
  426.                          patch_size=4, embed_dims=[96, 192, 384, 768], num_heads=[4, 8, 16, 32],
  427.                          mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
  428.                          norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[5, 5, 23, 5], sr_ratios=[8, 4, 2, 1],
  429.                          drop_rate=0.0, drop_path_rate=0.6,)
  432. if __name__=='__main__':
  433.     model = transnext_micro(factor=0.5, depth_factor=0.5)
  434.     model.eval()
  435.     inputs = torch.randn(1,3,640,640)
  436.     for i in model(inputs):
  437.         print(i.size())

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

4.1 修改一

第一步还是建立文件，我们找到如下ultralytics/nn文件夹下建立一个目录名字呢就是'Addmodules'文件夹( 用群内的文件的话已经有了无需新建) ！然后在其内部建立一个新的py文件将核心代码复制粘贴进去即可

4.2 修改二

第二步我们在该目录下创建一个新的py文件名字为'__init__.py'( 用群内的文件的话已经有了无需新建) ，然后在其内部导入我们的检测头如下图所示。

4.3 修改三

第三步我门中到如下文件'ultralytics/nn/tasks.py'进行导入和注册我们的模块( 用群内的文件的话已经有了无需重新导入直接开始第四步即可) ！

从今天开始以后的教程就都统一成这个样子了，因为我默认大家用了我群内的文件来进行修改！！

4.4 修改四

添加如下两行代码！！！

​

4.5 修改五

找到七百多行大概把具体看图片，按照图片来修改就行，添加红框内的部分，注意没有()只是 函数 名。 

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

4.6 修改六

下面的两个红框内都是需要改动的。

​ 

  1.         if isinstance(c2, list):
  2.             m_ = m
  3.             m_.backbone = True
  4.         else:
  5.             m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
  6.             t = str(m)[8:-2].replace('__main__.', '')  # module type
  9.         m.np = sum(x.numel() for x in m_.parameters())  # number params
  10.         m_.i, m_.f, m_.type = i + 4 if backbone else i, f, t  # attach index, 'from' index, type

4.7 修改七

如下的也需要修改，全部按照我的来。

​

代码如下把原先的代码替换了即可。 

  1.         if verbose:
  2.             LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f}  {t:<45}{str(args):<30}')  # print
  4.         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
  5.         layers.append(m_)
  6.         if i == 0:
  7.             ch = []
  8.         if isinstance(c2, list):
  9.             ch.extend(c2)
  10.             if len(c2) != 5:
  11.                 ch.insert(0, 0)
  12.         else:
  13.             ch.append(c2)

4.8 修改八

修改八和前面的都不太一样，需要修改前向传播中的一个部分， 已经离开了parse_model方法了。

可以在图片中开代码行数，没有离开task.py文件都是同一个文件。 同时这个部分有好几个前向传播都很相似，大家不要看错了， 是70多行左右的！！！，同时我后面提供了代码，大家直接复制粘贴即可，有时间我针对这里会出一个视频。

​​

代码如下-> 

  1.     def _predict_once(self, x, profile=False, visualize=False, embed=None):
  2.         """
  3.         Perform a forward pass through the network.
  5.         Args:
  6.             x (torch.Tensor): The input tensor to the model.
  7.             profile (bool):  Print the computation time of each layer if True, defaults to False.
  8.             visualize (bool): Save the feature maps of the model if True, defaults to False.
  9.             embed (list, optional): A list of feature vectors/embeddings to return.
  11.         Returns:
  12.             (torch.Tensor): The last output of the model.
  13.         """
  14.         y, dt, embeddings = [], [], []  # outputs
  15.         for m in self.model:
  16.             if m.f != -1:  # if not from previous layer
  17.                 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
  18.             if profile:
  19.                 self._profile_one_layer(m, x, dt)
  20.             if hasattr(m, 'backbone'):
  21.                 x = m(x)
  22.                 if len(x) != 5:  # 0 - 5
  23.                     x.insert(0, None)
  24.                 for index, i in enumerate(x):
  25.                     if index in self.save:
  26.                         y.append(i)
  27.                     else:
  28.                         y.append(None)
  29.                 x = x[-1]  # 最后一个输出传给下一层
  30.             else:
  31.                 x = m(x)  # run
  32.                 y.append(x if m.i in self.save else None)  # save output
  33.             if visualize:
  34.                 feature_visualization(x, m.type, m.i, save_dir=visualize)
  35.             if embed and m.i in embed:
  36.                 embeddings.append(nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1))  # flatten
  37.                 if m.i == max(embed):
  38.                     return torch.unbind(torch.cat(embeddings, 1), dim=0)
  39.         return x

到这里就完成了修改部分，但是这里面细节很多，大家千万要注意不要替换多余的代码，导致报错，也不要拉下任何一部，都会导致运行失败，而且报错很难排查！！！很难排查！！！

注意！！！ 额外的修改（本文需要修改）！

关注我的其实都知道，我大部分的修改都是一样的，这个网络需要额外的修改一步，就是s一个参数，将下面的s改为640！！！即可完美运行！！

打印计算量问题解决方案

我们找到如下文件'ultralytics/utils/torch_utils.py'按照如下的图片进行修改，否则容易打印不出来计算量。

​

注意事项（本文需要修改）！！！

如果大家在验证的时候报错形状不匹配的错误可以固定验证集的图片尺寸，方法如下 ->

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

五、TransNeXt的yaml文件

复制如下yaml文件进行运行！！！

5.1 TransNeXt 的yaml文件版本1

此版本训练信息：YOLO11-TransNeXt summary: 498 layers, 2,029,091 parameters, 2,029,075 gradients, 6.0 GFLOPs

使用说明：# 下面 [-1, 1,transnext_micro, [0.25，0.5]] 参数位置的0.25是通道放缩的系数, YOLOv11N是0.25 YOLOv11S是0.5 YOLOv11M是1. YOLOv11l是1 YOLOv11是1.5大家根据自己训练的YOLO版本设定即可.
#  0.5对应的是模型的深度系数

# # 支持的版本:__all__ = ['transnext_micro', 'transnext_small', 'transnext_base', 'transnext_tiny'] 

  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  4. # Parameters
  5. nc: 80 # number of classes
  6. scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  7.   # [depth, width, max_channels]
  8.   n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  9.   s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  10.   m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  11.   l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  12.   x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
  14. # 下面 [-1, 1,transnext_micro, [0.25，0.5]] 参数位置的0.25是通道放缩的系数, YOLOv11N是0.25 YOLOv11S是0.5 YOLOv11M是1. YOLOv11l是1 YOLOv11是1.5大家根据自己训练的YOLO版本设定即可.
  15.                                 #  0.5对应的是模型的深度系数
  16. # 支持的版本:__all__ = ['transnext_micro', 'transnext_small', 'transnext_base', 'transnext_tiny']
  17. # YOLO11n backbone
  18. backbone:
  19.   # [from, repeats, module, args]
  20.   - [-1, 1, transnext_micro, [0.25, 0.5]] # 0-4 P1/2
  21.   - [-1, 1, SPPF, [1024, 5]] # 5
  22.   - [-1, 2, C2PSA, [1024]] # 6
  24. # YOLO11n head
  25. head:
  26.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  27.   - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  28.   - [-1, 2, C3k2, [512, False]] # 9
  30.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  31.   - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  32.   - [-1, 2, C3k2, [256, False]] # 12 (P3/8-small)
  34.   - [-1, 1, Conv, [256, 3, 2]]
  35.   - [[-1, 9], 1, Concat, [1]] # cat head P4
  36.   - [-1, 2, C3k2, [512, False]] # 15 (P4/16-medium)
  38.   - [-1, 1, Conv, [512, 3, 2]]
  39.   - [[-1, 6], 1, Concat, [1]] # cat head P5
  40.   - [-1, 2, C3k2, [1024, True]] # 18 (P5/32-large)
  42.   - [[12, 15, 18], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.2 训练文件 

  1. import warnings
  2. warnings.filterwarnings('ignore')
  3. from ultralytics import YOLO
  5. if __name__ == '__main__':
  6.     model = YOLO('ultralytics/cfg/models/v8/yolov8-C2f-FasterBlock.yaml')
  7.     # model.load('yolov8n.pt') # loading pretrain weights
  8.     model.train(data=r'替换数据集yaml文件地址',
  9.                 # 如果大家任务是其它的'ultralytics/cfg/default.yaml'找到这里修改task可以改成detect, segment, classify, pose
  10.                 cache=False,
  11.                 imgsz=640,
  12.                 epochs=150,
  13.                 single_cls=False,  # 是否是单类别检测
  14.                 batch=4,
  15.                 close_mosaic=10,
  16.                 workers=0,
  17.                 device='0',
  18.                 optimizer='SGD', # using SGD
  19.                 # resume='', # 如过想续训就设置last.pt的地址
  20.                 amp=False,  # 如果出现训练损失为Nan可以关闭amp
  21.                 project='runs/train',
  22.                 name='exp',
  23.                 )

六、成功运行记录

下面是成功运行的截图，已经完成了有1个epochs的训练，图片太大截不全第2个epochs，这里改完之后打印出了点问题，但是不影响任何功能，后期我找时间修复一下这个问题。

​​

七、本文总结

到此本文的正式分享内容就结束了，在这里给大家推荐我的YOLOv11改进有效涨点专栏，本专栏目前为新开的平均质量分98分，后期我会根据各种最新的前沿顶会进行论文复现，也会对一些老的改进机制进行补充， 目前本专栏免费阅读(暂时，大家尽早关注不迷路~) ，如果大家觉得本文帮助到你了，订阅本专栏，关注后续更多的更新~

​​