yolov26改进 | 主干/Backboen篇 | 目标检测网络最新成果TransNeXt特征提取网络(yolov26全系列可用轻量化)

开始正文前,先向大家推荐我的YOLO专栏系列。本人持续更新 YOLOv8、YOLO11、YOLO26 等热门模型,内容覆盖图像分类、目标检测、实例分割、多目标跟踪、姿态估计与关键点检测,重点讲解 小目标检测、注意力机制、特征融合、损失函数改进、自定义数据集训练、消融实验及论文代码复现。同时分享如何使用 OpenAI Codex辅助撰写论文、配置实验环境、调试项目、改进模型和分析实验结果。 专栏目前正在进行限时优惠,每周更新 5–7篇最新论文机制、YOLO改进方法和实战教程。订阅后可获得包含本人全部改进方案的代码与配置文件,并加入专属技术交流群。我也会定期在群内分享 YOLO论文选题、创新点设计、实验方案、论文写作与投稿发表经验,欢迎大家订阅交流!


一、本文介绍

本文给大家带来的改进机制是视觉特征提取网络——TransNeXt,本文将其应用到YOLOv26的特征提取部分,利用其更强的全局建模能力和信息交互能力,帮助网络提取更加丰富、有效的多尺度特征。TransNeXt的核心结构主要包括聚合像素聚焦注意力机制卷积GLU模块,其通过模拟生物视觉系统中的中心凹感知方式,使每个像素在重点关注局部区域信息的同时,也能够与全局特征进行交互,从而兼顾局部细节提取和长距离依赖建模,卷积GLU则进一步增强了空间信息与通道信息的混合能力,使网络能够更加充分地学习目标的纹理、边缘和上下文特征。TransNeXt在图像分类、目标检测和语义分割等多种视觉任务中均展现出了较强的性能,尤其适合复杂背景和多尺度目标检测任务。需要注意的是,该网络结构相对复杂,训练时间通常较长,对显存和计算资源也有一定要求,大家在使用时需要根据自己的硬件条件选择合适的模型版本。同时,本文还会解决TransNeXt集成过程中容易出现的相关报错,并结合完整代码和配置文件,手把手讲解如何将该网络结构添加到YOLOv26中。

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

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

专栏链接:YOLOv26有效涨点专栏包含:Conv、注意力机制、主干/Backbone、损失函数、优化器、后处理等改进机制


目录

一、本文介绍

二、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 修改九

五、TransNeXt的yaml文件

5.1 TransNeXt的yaml文件

5.2 训练文件的代码

六、成功运行记录 

七、本文总结


二、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的核心代码 

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

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

__all__ = ['transnext_micro', 'transnext_small', 'transnext_base', 'transnext_tiny']

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, factor=0.5, depth_factor=0.5, 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__()
        embed_dims = tuple([int(dim * factor) for dim in embed_dims])
        depths = [max(round(dim * depth_factor), 1) if dim > 1 else dim for dim in depths]
        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, factor, depth_factor, **kwargs):
        super().__init__(factor, depth_factor, 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,factor, depth_factor,  **kwargs):
        super().__init__(factor, depth_factor, 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, factor, depth_factor, **kwargs):
        super().__init__(factor, depth_factor, 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,factor, depth_factor,  **kwargs):
        super().__init__(factor, depth_factor, 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,)


if __name__=='__main__':
    model = transnext_micro(factor=0.5, depth_factor=0.5)
    model.eval()
    inputs = torch.randn(1,3,640,640)
    for i in model(inputs):
        print(i.size())


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

 4.1 修改一

我们复制网络结构代码到“ultralytics/nn”目录下创建一个py文件复制粘贴进去 。


4.2 修改二 

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


4.3 修改三 

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

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


4.4 修改四 

添加如下两行代码!!!


4.5 修改五

找到1600多行大概把具体看图片,按照图片来修改就行,添加红框内的部分,注意没有()只是函数名,我这里只添加了部分的版本,这个还有更多的版本可以添加可以看后面章节五,看我给的代码函数头即可。

        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

4.7 修改七

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

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

        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.8 修改八

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

可以在图片中看代码行数,没有离开task.py文件都是同一个文件。 同时这个部分有好几个前向传播都很相似,大家不要看错了,是160多行左右的(不同仓库版本可能有些差异)!!!,同时我后面提供了代码,大家直接复制粘贴即可,不会修改联系博主获取视频教程。

​​

代码如下->

    def _predict_once(self, x, profile=False, visualize=False, embed=None):
        """Perform a forward pass through the network.

        Args:
            x (torch.Tensor): The input tensor to the model.
            profile (bool): Print the computation time of each layer if True.
            visualize (bool): Save the feature maps of the model if True.
            embed (list, optional): A list of layer indices to return embeddings from.

        Returns:
            (torch.Tensor): The last output of the model.
        """
        y, dt, embeddings = [], [], []  # outputs
        embed = frozenset(embed) if embed is not None else {-1}
        max_idx = max(embed)
        for m in self.model:
            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_idx:
                    return torch.unbind(torch.cat(embeddings, 1), dim=0)
        return x

4.9 修改九

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


五、TransNeXt的yaml文件

5.1 TransNeXt的yaml文件

训练信息:YOLO26-Backbone-TransNeXt summary: 435 layers, 1,940,516 parameters, 1,940,516 gradients, 5.4 GFLOPs

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# Ultralytics YOLO26 object detection model with P3/8 - P5/32 outputs
# Model docs: https://docs.ultralytics.com/models/yolo26
# Task docs: https://docs.ultralytics.com/tasks/detect

# Parameters
nc: 80 # number of classes
end2end: True # whether to use end-to-end mode
reg_max: 1 # DFL bins
scales: # model compound scaling constants, i.e. 'model=yolo26n.yaml' will call yolo26.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.50, 0.25, 1024] # summary: 260 layers, 2,572,280 parameters, 2,572,280 gradients, 6.1 GFLOPs
  s: [0.50, 0.50, 1024] # summary: 260 layers, 10,009,784 parameters, 10,009,784 gradients, 22.8 GFLOPs
  m: [0.50, 1.00, 512] # summary: 280 layers, 21,896,248 parameters, 21,896,248 gradients, 75.4 GFLOPs
  l: [1.00, 1.00, 512] # summary: 392 layers, 26,299,704 parameters, 26,299,704 gradients, 93.8 GFLOPs
  x: [1.00, 1.50, 512] # summary: 392 layers, 58,993,368 parameters, 58,993,368 gradients, 209.5 GFLOPs

# 我提供了版本分别是对应是 ['transnext_micro, transnext_small, transnext_base, transnext_tiny']

# YOLO26 backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, transnext_micro, [0.25, 0.50]] # 0-4 P1/2 这里是四层大家不要被yaml文件限制住了思维,不会画图进群看视频.
  # 注意args位置的参数对应模型的通道放缩系数width在上面scales位置, 假设你用yolov26n那么可以设置0.25 如果你用yolov26s可以设置0.5
  # 0.5对应的是模型的深度系数
  - [-1, 1, SPPF, [1024, 5, 3, True]] # 5
  - [-1, 2, C2PSA, [1024]] # 6

# YOLO26 head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, True]] # 9

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, True]] # 12 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 9], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, True]] # 15 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 6], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True, 0.5, True]] # 18 (P5/32-large)

  - [[12, 15, 18], 1, Detect, [nc]] # Detect(P3, P4, P5)


5.2 训练文件的代码

可以复制我的运行文件进行运行。

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('yolov8-MLLA.yaml')
    # 如何切换模型版本, 上面的ymal文件可以改为 yolov8s.yaml就是使用的v8s,
    # 类似某个改进的yaml文件名称为yolov8-XXX.yaml那么如果想使用其它版本就把上面的名称改为yolov8l-XXX.yaml即可(改的是上面YOLO中间的名字不是配置文件的)!
    # model.load('yolov8n.pt') # 是否加载预训练权重,科研不建议大家加载否则很难提升精度
    model.train(data=r"C:\Users\Administrator\PycharmProjects\yolov5-master\yolov5-master\Construction Site Safety.v30-raw-images_latestversion.yolov8\data.yaml",
                # 如果大家任务是其它的'ultralytics/cfg/default.yaml'找到这里修改task可以改成detect, segment, classify, pose
                cache=False,
                imgsz=640,
                epochs=150,
                single_cls=False,  # 是否是单类别检测
                batch=16,
                close_mosaic=0,
                workers=0,
                device='0',
                optimizer='SGD', # using SGD
                # resume='runs/train/exp21/weights/last.pt', # 如过想续训就设置last.pt的地址
                amp=True,  # 如果出现训练损失为Nan可以关闭amp
                project='runs/train',
                name='exp',
                )


六、成功运行记录 

下面是成功运行的截图,已经完成了有1个epochs的训练,图片太大截不全第2个epochs了。 


七、本文总结

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

专栏链接:YOLOv26有效涨点专栏包含:Conv、注意力机制、主干/Backbone、损失函数、优化器、后处理等改进机制

​​

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Snu77

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值