TransNeXt 深度解析:聚合注意力机制的突破性视觉骨干网络

TransNeXt 深度解析:聚合注意力机制的突破性视觉骨干网络

目录

  1. 引言
  2. 背景与动机
  3. 核心创新点
  4. 网络架构详解
  5. 关键组件源码解析
  6. 模型变体与配置
  7. 在YOLO11中的集成应用
  8. 实验结果与性能对比
  9. 使用示例
  10. 总结与展望

引言

TransNeXt 是 2024 年提出的一种创新性视觉骨干网络,论文标题为 "TransNeXt: Robust Foveal Visual Perception for Vision Transformers" 。该模型的核心创新是**聚合注意力(Aggregated Attention)机制,它巧妙地结合了局部窗口注意力和全局池化注意力,模拟人类视觉系统的中央凹视觉感知(Foveal Visual Perception)**特性。

论文信息

  • 论文标题: TransNeXt: Robust Foveal Visual Perception for Vision Transformers
  • 作者: Dai Shi
  • 发表时间: 2024
  • 代码: 提供 CUDA 加速版本和纯 PyTorch Native 版本

背景与动机

Vision Transformer 的局限性

传统的 Vision Transformer (ViT) 使用全局自注意力机制:

  • 计算复杂度高:O(N²),N 为 token 数量
  • 对位置编码敏感:绝对位置编码导致分辨率泛化能力差
  • 缺乏局部归纳偏置:需要大量数据才能学习到有效的局部特征

现有解决方案的问题

  1. 窗口注意力(如 Swin Transformer)

    • 仅在局部窗口内计算注意力
    • 需要窗口移位操作增加复杂度
    • 跨窗口信息交互受限
  2. 池化注意力(如 PVT)

    • 通过下采样减少 Key/Value 的数量
    • 会丢失局部细节信息

TransNeXt 的解决思路

受人类视觉系统启发,TransNeXt 提出聚合注意力

  • 中央凹区域(局部):高分辨率,关注细节
  • 周围区域(全局):低分辨率,获取上下文

这种设计既保留了局部细节,又能捕获全局上下文,同时保持线性计算复杂度。


核心创新点

1. 聚合注意力(Aggregated Attention)

聚合注意力是 TransNeXt 的核心,它将局部和全局注意力统一到一个 Softmax 中:

复制代码
Query 同时与:
  ├── 局部邻域的 Key/Value(unfold 操作获取窗口)
  └── 全局池化的 Key/Value(自适应平均池化)
进行注意力计算,然后统一 Softmax 归一化

关键优势

  • 局部和全局信息在同一个注意力权重空间中竞争
  • 自然实现了信息的层次化聚合
  • 线性计算复杂度 O(N)

2. 像素聚焦注意力(Pixel-focused Attention)

为解决传统注意力机制的"token 聚焦"问题,TransNeXt 引入:

python 复制代码
# 传统注意力:只关注特定 token
attn = softmax(Q @ K.T / sqrt(d)) @ V

# 像素聚焦注意力:增加可学习的 token-to-token 交互
attn_local = (q_norm @ learnable_tokens) + learnable_bias + attn_local

通过引入可学习的 token 交互矩阵,每个像素位置可以动态调整其注意力分布。

3. 序列长度缩放的余弦注意力

TransNeXt 使用改进的注意力计算方式:

python 复制代码
# L2 归一化 Query 和 Key
q_norm = F.normalize(q, dim=-1)
k_norm = F.normalize(k, dim=-1)

# 添加 Query Embedding 并使用可学习温度
attn = (q_norm + query_embedding) * softplus(temperature) * seq_length_scale @ k_norm.T

设计要点

  • L2 归一化:稳定训练,防止注意力塌缩
  • Query Embedding:增加 query 的表达能力
  • 可学习温度:使用 softplus 确保温度非负
  • 序列长度缩放:log(N) 缩放因子,适应不同分辨率

4. 连续相对位置偏置(Continuous Positional Bias, CPB)

使用 MLP 生成连续的相对位置偏置,支持任意分辨率:

python 复制代码
# MLP 生成位置偏置
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_coords = sign(coords) * log2(abs(coords) + 1) / log2(8)

5. 卷积门控线性单元(Convolutional GLU)

TransNeXt 使用改进的 FFN,引入门控机制和深度卷积:

python 复制代码
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

网络架构详解

整体架构

TransNeXt 采用四阶段金字塔结构:

复制代码
输入图像 (H, W, 3)
    │
    ▼
┌──────────────────────┐
│  Patch Embedding 1   │  7×7 Conv, stride=4
│  (embed_dims[0])     │  输出: H/4 × W/4
└──────────┬───────────┘
           │
    ▼ Stage 1 ▼
┌──────────────────────┐
│  N × TransNeXt Block │  AggregatedAttention + ConvGLU
│  (embed_dims[0])     │  sr_ratio=8
└──────────┬───────────┘
           │
    ▼ Patch Embedding ▼   3×3 Conv, stride=2
           │
    ▼ Stage 2 ▼
┌──────────────────────┐
│  N × TransNeXt Block │  sr_ratio=4
│  (embed_dims[1])     │
└──────────┬───────────┘
           │
    ▼ Patch Embedding ▼
           │
    ▼ Stage 3 ▼
┌──────────────────────┐
│  N × TransNeXt Block │  sr_ratio=2 (最深阶段)
│  (embed_dims[2])     │
└──────────┬───────────┘
           │
    ▼ Patch Embedding ▼
           │
    ▼ Stage 4 ▼
┌──────────────────────┐
│  N × TransNeXt Block │  sr_ratio=1 (全局注意力)
│  (embed_dims[3])     │  使用 Attention 而非 AggregatedAttention
└──────────────────────┘
           │
    ▼ 输出 ▼
  [P2, P3, P4, P5]  四个尺度的特征图

TransNeXt Block 结构

复制代码
输入 x (B, N, C)
    │
    ├───────────────────────┐
    ▼                       │ (残差连接)
┌──────────────────────┐    │
│     LayerNorm        │    │
└──────────┬───────────┘    │
           │                │
┌──────────────────────┐    │
│  AggregatedAttention │    │  或 Attention (sr_ratio=1)
│  (局部 + 全局)       │    │
└──────────┬───────────┘    │
           │                │
┌──────────────────────┐    │
│     DropPath         │    │
└──────────┬───────────┘    │
           + ◄──────────────┘
           │
    ├───────────────────────┐
    ▼                       │ (残差连接)
┌──────────────────────┐    │
│     LayerNorm        │    │
└──────────┬───────────┘    │
           │                │
┌──────────────────────┐    │
│  ConvolutionalGLU    │    │  门控 + 深度卷积
└──────────┬───────────┘    │
           │                │
┌──────────────────────┐    │
│     DropPath         │    │
└──────────┬───────────┘    │
           + ◄──────────────┘
           │
    ▼ 输出 ▼

关键组件源码解析

1. 聚合注意力(AggregatedAttention)

python 复制代码
class AggregatedAttention(nn.Module):
    """聚合注意力 - TransNeXt 的核心创新"""
    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__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.sr_ratio = sr_ratio
        self.window_size = window_size
        self.local_len = window_size ** 2  # 局部窗口的 token 数量

        # 池化后的特征图尺寸
        self.pool_H = input_resolution[0] // sr_ratio
        self.pool_W = input_resolution[1] // 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)
        )

        # Query/Key/Value 投影
        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)

        # Query Embedding
        self.query_embedding = nn.Parameter(
            nn.init.trunc_normal_(torch.empty(num_heads, 1, self.head_dim),
                                  mean=0, std=0.02)
        )

        # 全局特征池化组件
        self.pool = nn.AdaptiveAvgPool2d((self.pool_H, self.pool_W))
        self.sr = nn.Conv2d(dim, dim, kernel_size=1)
        self.norm = nn.LayerNorm(dim)
        self.act = nn.GELU()

        # 连续位置偏置 MLP
        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)

        # 局部位置偏置
        self.relative_pos_bias_local = nn.Parameter(
            nn.init.trunc_normal_(torch.empty(num_heads, self.local_len),
                                  mean=0, std=0.0004)
        )

        # 动态局部偏置
        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))

2. 聚合注意力的前向传播

python 复制代码
def forward(self, x, H, W, relative_pos_index, relative_coords_table):
    B, N, C = x.shape

    # ========== 1. 生成 Query ==========
    # L2 归一化 + Query Embedding + 温度缩放
    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

    # ========== 2. 局部 Key/Value(Unfold 操作)==========
    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)
    k_local = k_local.reshape(B, N, -1)
    kv_local = torch.cat([k_local, v_local], dim=-1).permute(0, 2, 1).reshape(B, -1, H, W)

    # Unfold 提取每个位置的局部邻域
    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)

    # 局部注意力矩阵
    attn_local = ((q_norm_scaled.unsqueeze(-2) @ k_local).squeeze(-2)
                  + self.relative_pos_bias_local.unsqueeze(1))
    attn_local = attn_local.masked_fill(self.padding_mask, float('-inf'))

    # ========== 3. 全局 Key/Value(池化操作)==========
    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_)

    kv_pool = self.kv(x_).reshape(B, self.pool_len, 2 * self.num_heads, self.head_dim)
    kv_pool = kv_pool.permute(0, 2, 1, 3)
    k_pool, v_pool = kv_pool.chunk(2, dim=1)

    # 全局位置偏置
    pool_bias = self.cpb_fc2(self.cpb_act(self.cpb_fc1(relative_coords_table)))
    pool_bias = pool_bias.transpose(0, 1)[:, relative_pos_index.view(-1)].view(-1, N, self.pool_len)

    # 全局注意力矩阵
    attn_pool = q_norm_scaled @ F.normalize(k_pool, dim=-1).transpose(-2, -1) + pool_bias

    # ========== 4. 统一 Softmax ==========
    attn = torch.cat([attn_local, attn_pool], dim=-1).softmax(dim=-1)
    attn = self.attn_drop(attn)

    # ========== 5. 分别聚合局部和全局 Value ==========
    attn_local, attn_pool = torch.split(attn, [self.local_len, self.pool_len], dim=-1)

    # 局部注意力 + 像素聚焦
    attn_local = (q_norm @ self.learnable_tokens) + self.learnable_bias + attn_local
    x_local = (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)
    x = self.proj(x)
    x = self.proj_drop(x)

    return x

3. 相对位置编码生成

python 复制代码
@torch.no_grad()
def get_relative_position_cpb(query_size, key_size, pretrain_size=None):
    """生成连续相对位置偏置的坐标表和索引"""
    pretrain_size = pretrain_size or query_size

    # 生成 Query 和 Key 的坐标轴
    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

模型变体与配置

TransNeXt 提供四种规模的模型:

变体 embed_dims num_heads depths 参数量 适用场景
Micro [48, 96, 192, 384] [2, 4, 8, 16] [2, 2, 15, 2] ~12.8M 轻量级部署
Tiny [72, 144, 288, 576] [3, 6, 12, 24] [2, 2, 15, 2] ~28.2M 标准基线
Small [72, 144, 288, 576] [3, 6, 12, 24] [5, 5, 22, 5] ~49.7M 高精度
Base [96, 192, 384, 768] [4, 8, 16, 32] [5, 5, 23, 5] ~89.7M 最高精度

创建模型函数

python 复制代码
def transnext_micro(pretrained=False, **kwargs):
    """创建 TransNeXt-Micro 模型"""
    model = TransNeXt(
        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],
        **kwargs
    )
    return model

def transnext_tiny(pretrained=False, **kwargs):
    """创建 TransNeXt-Tiny 模型"""
    model = TransNeXt(
        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],
        **kwargs
    )
    return model

在YOLO11中的集成应用

YAML 配置文件

yaml 复制代码
# yolo11-TransNext.yaml
nc: 80  # 类别数

backbone:
  - [-1, 1, transnext_micro, []]  # TransNeXt 骨干
  - [-1, 1, SPPF, [1024, 5]]       # 空间金字塔池化
  - [-1, 2, C2PSA, [1024]]         # 注意力模块

head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 3], 1, Concat, [1]]      # 拼接 P4
  - [-1, 2, C3k2, [512, False]]

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 2], 1, Concat, [1]]      # 拼接 P3
  - [-1, 2, C3k2, [256, False]]

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 9], 1, Concat, [1]]
  - [-1, 2, C3k2, [512, False]]

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 6], 1, Concat, [1]]
  - [-1, 2, C3k2, [1024, True]]

  - [[12, 15, 18], 1, Detect, [nc]]

CUDA 加速与 Native 版本

TransNeXt 提供两个实现版本:

python 复制代码
# ultralytics/nn/backbone/TransNext.py
try:
    import swattention  # CUDA 加速库
    from ultralytics.nn.backbone.TransNeXt.TransNext_cuda import *
except ImportError as e:
    # 回退到纯 PyTorch 实现
    from ultralytics.nn.backbone.TransNeXt.TransNext_native import *
    pass

CUDA 版本优势

  • 使用 swattention 库加速滑动窗口注意力
  • 适合大规模训练和推理

Native 版本优势

  • 无需额外 CUDA 依赖
  • 更好的跨平台兼容性

实验结果与性能对比

ImageNet 分类性能

模型 参数量 GFLOPs Top-1 Acc (%)
Swin-T 29M 4.5 81.3
PVT-M 44M 6.7 81.2
TransNeXt-T 28M 5.7 84.0
Swin-S 50M 8.7 83.0
TransNeXt-S 50M 10.3 84.7

目标检测性能(COCO)

Backbone AP AP50 AP75
Swin-T 46.0 68.1 50.3
PVT-M 44.2 66.0 48.2
TransNeXt-T 47.9 69.5 52.4

使用示例

1. 基础使用

python 复制代码
from ultralytics import YOLO

# 加载 TransNeXt 配置
model = YOLO('ultralytics/cfg/models/11/yolo11-TransNext.yaml')

# 训练模型
results = model.train(
    data='coco.yaml',
    epochs=100,
    imgsz=640,
    batch=16,
    device=0
)

2. 单独使用骨干网络

python 复制代码
import torch
from ultralytics.nn.backbone.TransNext import transnext_micro

# 创建模型
model = transnext_micro()

# 测试前向传播
x = torch.randn(1, 3, 640, 640)
features = model(x)
for i, feat in enumerate(features):
    print(f"Stage {i+1}: {feat.shape}")
# 输出:
# Stage 1: torch.Size([1, 48, 160, 160])   - P2
# Stage 2: torch.Size([1, 96, 80, 80])     - P3
# Stage 3: torch.Size([1, 192, 40, 40])    - P4
# Stage 4: torch.Size([1, 384, 20, 20])    - P5

3. 自定义配置

python 复制代码
from ultralytics.nn.backbone.TransNeXt.TransNext_native import TransNeXt
from functools import partial
import torch.nn as nn

# 自定义 TransNeXt
model = TransNeXt(
    img_size=640,              # 输入分辨率
    pretrain_size=224,         # 预训练分辨率
    window_size=[3, 3, 3, None],
    patch_size=4,
    embed_dims=[64, 128, 256, 512],  # 自定义通道数
    num_heads=[2, 4, 8, 16],
    mlp_ratios=[8, 8, 4, 4],
    depths=[2, 2, 10, 2],      # 自定义深度
    sr_ratios=[8, 4, 2, 1],
    drop_path_rate=0.1,        # Stochastic Depth
    qkv_bias=True,
    norm_layer=partial(nn.LayerNorm, eps=1e-6)
)

总结与展望

TransNeXt 的主要贡献

  1. 聚合注意力:优雅地统一局部和全局注意力,模拟中央凹视觉感知
  2. 像素聚焦注意力:解决 token 聚焦问题,增强局部信息保留
  3. 序列长度缩放:使模型能够适应不同输入分辨率
  4. 连续位置偏置:支持任意分辨率的位置编码

适用场景

  • 高分辨率图像:聚合注意力在高分辨率下保持线性复杂度
  • 密集预测任务:像素聚焦机制增强定位精度
  • 多尺度目标:四阶段金字塔结构处理不同尺度
  • 实时推理:Micro 版本适合边缘部署

与其他模型对比

特性 TransNeXt Swin Transformer PVT
局部注意力 ✓(窗口 Unfold) ✓(移位窗口)
全局注意力 ✓(池化) ✓(下采样 K/V)
统一 Softmax
位置编码 CPB(连续) RPB(离散) APE
计算复杂度 O(N) O(N) O(N)

未来发展方向

  1. 更高效的窗口注意力 CUDA 实现
  2. 与自监督预训练结合(如 MAE)
  3. 多模态扩展(图像+语言)

参考资料


致谢

感谢 Ultralytics 团队为 YOLOv11 提供了如此强大的基础框架,使我们能够专注于模型的创新部分。同时,也要感谢 Dai Shi 研究团队分享了 TransNeXt 的优秀工作,为计算机视觉领域带来了新的突破。

相关推荐
AI小云1 小时前
【数据操作与可视化】Matplotlib绘图-基础功能
python·数据可视化
iiiiii111 小时前
【论文阅读笔记】IDAQ:离线元强化学习中的分布内在线适应
论文阅读·人工智能·笔记·学习·算法·机器学习·强化学习
Cleaner1 小时前
大模型的手和脚:从提示工程到 MCP
人工智能·llm·mcp
极客BIM工作室1 小时前
Sora模型双路径压缩网络详解
人工智能·深度学习
m***66731 小时前
Python 爬虫实战案例 - 获取社交平台事件热度并进行影响分析
开发语言·爬虫·python
DeepFlow 零侵扰全栈可观测1 小时前
DeepFlow 全栈可观测性 护航某银行核心系统全生命周期
数据库·人工智能·分布式·云原生·金融
I_ltt_Itw,1 小时前
Python多线程
服务器·python·pycharm
Hello.Reader1 小时前
Flink SQL 数据类型从 INT 到 VARIANT 的完整实战指南
python·sql·flink
哈哈哈也不行吗1 小时前
从入门到精通:大角几何在教学中的深度应用
人工智能·几何画板·几何绘图·大角几何·数学绘图工具