YOLOv8改进有效涨点系列-＞适合多种检测场景的BiFormer注意力机制(Bi-level Routing Attention)

鱼弦

于 2024-10-08 12:00:00 发布

阅读量715

点赞数 12

分类专栏：人工智能时代文章标签： YOLO

本文链接：https://blog.csdn.net/feng1790291543/article/details/139786758

版权

人工智能时代专栏收录该内容

2 篇文章 1 订阅

订阅专栏

鱼弦：公众号【红尘灯塔】，CSDN博客专家、内容合伙人、新星导师、全栈领域优质创作者、51CTO(Top红人+专家博主) 、github开源爱好者（go-zero源码二次开发、游戏后端架构 https://github.com/Peakchen）

1. 简介

BiFormer注意力机制是一种近年来兴起的注意力机制，其核心思想是通过双层路由机制来提升注意力机制的性能。本文将介绍如何在YOLOv8模型结构中添加BiFormer注意力机制，以实现极限涨点。

2. BiFormer注意力机制原理详解

BiFormer注意力机制主要包括以下两个阶段：

第一阶段：空间特征路由

在第一阶段，BiFormer注意力机制将输入特征映射到更高维度的空间，并通过一个卷积层进行特征提取。然后，使用一个全局自注意力层对特征进行聚合，并通过一个MLP层进行非线性变换。最后，将聚合后的特征与输入特征进行残差连接。

第二阶段：通道特征路由

在第二阶段，BiFormer注意力机制将第一阶段的输出特征映射到更高维度的通道，并通过一个卷积层进行特征提取。然后，使用一个通道自注意力层对特征进行聚合，并通过一个MLP层进行非线性变换。最后，将聚合后的特征与第一阶段的输出特征进行残差连接。

BiFormer注意力机制通过双层路由机制，可以更好地提取特征之间的依赖关系，从而提升模型的性能。

3. BiFormer注意力机制应用场景解释

BiFormer注意力机制可以应用于各种计算机视觉任务，例如目标检测、图像分类、语义分割等。在目标检测任务中，BiFormer注意力机制可以用于提取不同尺度的特征信息，从而提升模型对不同大小目标的检测能力。

4. BiFormer注意力机制文献材料链接

Bi-level Routing Attention Network for Visual Recognition

5. BiFormer注意力机制应用示例产品

目前，BiFormer注意力机制已经应用于一些目标检测产品中，例如旷视科技的Megvii Face++、商汤科技的SenseTime FacePro等。

6. YOLOv8模型改进

在YOLOv8模型中，我们可以将BiFormer注意力机制替换部分标准注意力机制，以提升模型的性能。例如，我们可以将Neck模块中的CSPNet注意力机制替换为BiFormer注意力机制。

修改后的Neck模块代码:

def _fpn(x):
    p5 = Conv2D(512, kernel_size=1, use_bias=False)(x)
    p5 = BatchNormalization()(p5)
    p5 = LeakyReLU()(p5)

    p4 = _upsample(p5)
    p4 = _csp_attn(p4, 256)

    p3 = _upsample(p4)
    p3 = _csp_attn(p3, 128)

    return p5, p4, p3

1. 算法实现

1.1 BiFormer注意力机制实现

BiFormer注意力机制主要包括以下两个阶段：

第一阶段：空间特征路由

def spatial_routing_attention(x):
    # 空间特征路由
    c = Conv2D(x.shape[-1] * 2, kernel_size=1, use_bias=False)(x)
    c = BatchNormalization()(c)
    c = LeakyReLU()(c)

    # 全局自注意力层
    q = c
    k = c
    v = c
    q = q[:, None, :, :, :]
    k = k.transpose(0, 1, 3, 2, 4)
    w = tf.matmul(q, k)
    w = tf.nn.softmax(w, axis=-1)
    out = tf.matmul(w, v)
    out = tf.squeeze(out, axis=1)

    # MLP层
    out = tf.keras.layers.Dense(x.shape[-1], activation='relu')(out)
    out = tf.keras.layers.Dense(x.shape[-1])(out)

    # 残差连接
    out = out + x
    return out

第二阶段：通道特征路由

def channel_routing_attention(x):
    # 通道特征路由
    c = Conv2D(x.shape[-1] * 2, kernel_size=1, use_bias=False)(x)
    c = BatchNormalization()(c)
    c = LeakyReLU()(c)

    # 通道自注意力层
    q = c.transpose(0, 2, 1, 3)
    k = c.transpose(0, 2, 1, 3)
    v = c
    w = tf.matmul(q, k)
    w = tf.nn.softmax(w, axis=-1)
    out = tf.matmul(w, v)
    out = out.transpose(0, 1, 3, 2)

    # MLP层
    out = tf.keras.layers.Dense(x.shape[-1], activation='relu')(out)
    out = tf.keras.layers.Dense(x.shape[-1])(out)

    # 残差连接
    out = out + x
    return out

1.2 YOLOv8模型集成BiFormer注意力机制

修改后的Neck模块代码:

def _fpn(x):
    p5 = Conv2D(512, kernel_size=1, use_bias=False)(x)
    p5 = BatchNormalization()(p5)
    p5 = LeakyReLU()(p5)

    p4 = _upsample(p5)
    p4 = _bi_routing_attn(p4, 256)

    p3 = _upsample(p4)
    p3 = _bi_routing_attn(p3, 128)

    return p5, p4, p3

2. 代码完整详细实现

import tensorflow as tf
from tensorflow.keras import layers

# 定义数据集类
class YOLOv8Dataset(tf.keras.data.Dataset):
    def __init__(self, train_images, train_targets, batch_size, shuffle=True):
        super().__init__()
        self.train_images = train_images
        self.train_targets = train_targets
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __getitem__(self, index):
        images, targets = self.train_images[index], self.train_targets[index]
        images = tf.cast(images, tf.float32)
        images = tf.image.resize(images, (640, 640))
        images = tf.image.normalize(images, mean=[0.5, 0.5, 0.5], stddev=[0.5, 0.5, 0.5])