首页 > 其他分享 >YOLO11改进-注意力-引入非局部注意力机制NonLocalBlockND 增强小目标、遮挡检测

YOLO11改进-注意力-引入非局部注意力机制NonLocalBlockND 增强小目标、遮挡检测

时间:2024-10-21 20:46:37浏览次数:6  
标签:NonLocalBlockND YOLO11 nn self channels inter 注意力 size

           

        本篇文章将介绍一个新的改进机制——NonLocalBlockND,并阐述如何将其应用于YOLOv11中,显著提升模型性能。首先,我们将解析NonLocalBlockND的工作原理,非局部注意力机制NonLocalBlockND通过在全局范围内捕捉特征图中所有位置的相互关系,提升模型性能。它利用三个分支(θ、ϕ、g)生成降维后的特征表示,计算相似性并进行加权,从而结合局部与全局信息。随后,我们会详细说明如何将该模块与YOLOv11相结合,改善模型对全局信息的捕捉能力,从而提升目标检测的精度。最后,展示代码实现细节及其使用方法,最终展现这一改进对目标检测效果的积极影响。

1. 非局部注意力机制NonLocalBlockND结构介绍       

        非局部注意力机制是一种有效捕捉输入特征图中所有位置之间关系的方法。它特别适合处理需要全局上下文信息的任务。与传统卷积神经网络(CNN)相比,非局部块不再局限于局部邻域的特征。

非局部注意力机制通过在全局范围内捕捉特征图中所有位置的相互关系来提升模型性能。该机制的核心是通过三个分支(θ、ϕ、g)生成降维后的特征表示,用于计算特征之间的相似性并进行加权。

  1. θ分支:将输入特征降维到 512 通道,用于生成查询向量(query)。
  2. ϕ分支:同样降维为 512 通道,用于生成键向量(key),用于与查询向量计算相似性。
  3. g分支:降维为 512 通道,生成输入特征的值向量(value),用于加权聚合。

接着通过 θ 和 ϕ 的相似性计算生成注意力权重矩阵,并将其与 g 分支的特征相乘,得到全局加权后的特征表示。最后,通过 1x1 卷积恢复通道数,并采用残差结构保留原始特征与加权后的特征,结合局部与全局信息,提升模型的感知能力。

2. YOLOv11与NonLocalBlockND的结合   

        在YOLOv11的骨干网络中,NonLocalBlockND模块可以放置在SPPF模块之前。这样,网络不仅能够捕捉输入图像的局部细节,还能获取全局上下文信息。通过融合局部和全局特征,YOLOv11在目标检测中的识别能力将得到显著增强。

3. NonLocalBlockND代码部分

import torch
from torch import nn
from torch.nn import functional as F


class NonLocalBlockND(nn.Module):
    def __init__(self, in_channels, inter_channels=None,  sub_sample=True, bn_layer=True):
        super(NonLocalBlockND, self).__init__()

        self.sub_sample = sub_sample  # 是否进行下采样

        self.in_channels = in_channels  # 输入通道数
        self.inter_channels = inter_channels  # 中间通道数

        # 如果未指定中间通道数,默认为输入通道数的一半
        if self.inter_channels is None:
            self.inter_channels = in_channels // 2
            if self.inter_channels == 0:
                self.inter_channels = 1


        # 定义 g、theta、phi 的卷积层
        self.g = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels,
                         kernel_size=1, stride=1, padding=0)

        # 定义 W 层,可选择是否使用批归一化
        if bn_layer:
            self.W = nn.Sequential(
                nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels,
                        kernel_size=1, stride=1, padding=0),
                nn.BatchNorm2d(self.in_channels)
            )
            # nn.init.constant(self.W[1].weight, 0)  # 初始化权重
            # nn.init.constant(self.W[1].bias, 0)    # 初始化偏置
        else:
            self.W = nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels,
                             kernel_size=1, stride=1, padding=0)
            # nn.init.constant(self.W.weight, 0)  # 初始化权重
            # nn.init.constant(self.W.bias, 0)    # 初始化偏置

        # 定义 theta 和 phi 的卷积层
        self.theta = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels,
                             kernel_size=1, stride=1, padding=0)
        self.phi = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels,
                           kernel_size=1, stride=1, padding=0)

        # 如果进行下采样,则对 g 和 phi 添加池化层
        if sub_sample:
            self.g = nn.Sequential(self.g, nn.MaxPool2d(kernel_size=(2, 2)))
            self.phi = nn.Sequential(self.phi, nn.MaxPool2d(kernel_size=(2, 2)))

    def forward(self, x):
        '''
        前向传播方法
        :param x: 输入张量,形状为 (b, c, t, h, w) (对于3D数据)
        :return: 输出张量,形状与输入相同
        '''
        batch_size = x.size(0)  # 获取批量大小

        # 计算 g(x)
        g_x = self.g(x).view(batch_size, self.inter_channels, -1)  # 变形为 (b, inter_channels, N)
        g_x = g_x.permute(0, 2, 1)  # 调整维度顺序

        # 计算 theta(x) 和 phi(x)
        theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)
        theta_x = theta_x.permute(0, 2, 1)  # 调整维度顺序
        phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)

        # 计算注意力权重
        f = torch.matmul(theta_x, phi_x)  # 矩阵乘法
        f_div_C = F.softmax(f, dim=-1)  # 归一化

        # 加权聚合
        y = torch.matmul(f_div_C, g_x)  # 使用权重对 g_x 加权
        y = y.permute(0, 2, 1).contiguous()  # 调整维度顺序
        y = y.view(batch_size, self.inter_channels, *x.size()[2:])  # 变形为 (b, inter_channels, t, h, w)

        # 融合输入和输出
        W_y = self.W(y)  # 通过 W 层
        z = W_y + x  # 残差连接

        return z  # 返回最终输出


if __name__ == '__main__':
    model = NonLocalBlockND(256)
    # 生成一个1x256x80x80的张量
    data = torch.randn(1, 256, 80, 80)
    updata = model.forward(data)
    print(updata.shape)  # 输出张量的形状

 4. 将NonLocalBlockND引入到YOLOv11中

第一: 将下面的核心代码复制到D:\bilibili\model\YOLO11\ultralytics-main\ultralytics\nn路径下,如下图所示。

第二:在task.py中导入BCBlock包

第三:在task.py中的模型配置部分下面代码

在SPPF模块之前添加NonLocalBlockND

elif m is NonLocalBlockND:
    args = [ch[f]]

第四:将模型配置文件复制到YOLOV11.YAMY文件中

# Ultralytics YOLO 

标签:NonLocalBlockND,YOLO11,nn,self,channels,inter,注意力,size
From: https://blog.csdn.net/qq_64693987/article/details/142926536

相关文章