YOLOv8改进心得:基于Python实现的DAttention(DAT)注意力机制提升目标检测性能

引言

随着深度学习的不断发展,YOLO(You Only Look Once)系列模型凭借其高效的目标检测能力,已经成为计算机视觉领域中不可或缺的工具。YOLOv8作为YOLO系列最新版本,不仅在模型精度和推理速度上有显著提升,同时还集成了更多优化策略,广泛应用于无人驾驶、智能监控、医疗影像分析等领域。

然而,尽管YOLOv8已经具有极高的性能,但在处理复杂场景、密集目标和小目标检测时,仍存在进一步优化的空间。注意力机制作为一种能够让模型更加高效地关注关键信息的技术,近年来被广泛应用于各类视觉任务中,其中的DAttention(DAT)机制,凭借其可变形注意力(Deformable Attention)的独特设计,在目标检测中展现出了强大的表现。

本文将深入探讨如何将DAT注意力机制应用到YOLOv8中,以进一步提升目标检测的精度与性能。我们将详细介绍DAT的工作原理、代码实现以及如何将其整合到YOLOv8模型中,提供清晰易懂的步骤和逐行注释,确保您能够顺利实现并应用这一技术。


论文地址:https://openaccess.thecvf.com/content/CVPR2022/papers/Xia_Vision_Transformer_With_Deformable_Attention_CVPR_2022_paper.pdf

代码地址:https://github.com/LeapLabTHU/DAT

一、DAT注意力机制简介

1.1 DAT的引入背景

传统的Transformer自注意力机制处理输入图像中的每个像素点,这在捕捉全局上下文信息时表现出色,但在处理高分辨率图像时,计算量往往非常巨大,极大影响了模型的推理速度和效率。为了解决这一问题,Deformable Attention(可变形注意力)应运而生。

DAT(Vision Transformer with Deformable Attention)通过引入可变形注意力机制,仅在图像的关键区域进行计算,减少了冗余信息的处理,极大地提高了模型的效率和性能。该机制允许模型动态地选择采样点,从而集中资源在关键信息上,适用于图像分类、目标检测等任务。

1.2 DAT的核心思想

DAT的核心思想可以概括为以下几点:

  1. 可变形注意力:与传统Transformer的全局自注意力不同,DAT通过动态选择采样点,只关注图像中的关键区域,减少了计算量。
  2. 动态采样点:DAT允许模型根据输入图像的特定区域自动调整采样点位置,使得注意力机制可以灵活地捕捉重要的特征。
  3. 即插即用设计:DAT机制的灵活设计使其可以无缝集成到不同的视觉任务中,如目标检测、图像分类等,极大提升了模型在多任务下的适应性。

1.3 DAT与其他注意力机制的对比

DAT与传统的自注意力机制和卷积神经网络中的可变形卷积(DCN)相比,最大的区别在于它的灵活性和计算效率。传统的自注意力机制需要计算全局信息,计算量随图像分辨率的增大而快速增长;而DAT通过动态调整采样点,减少了无关区域的计算,显著提升了效率。

此外,与DCN不同,DAT可以同时处理不同的图像内容和大小,具有更广泛的适用性。相比之下,DCN更多地用于局部特征的提取,而DAT则能够结合全局上下文信息,从而在处理复杂场景时表现更为出色。


二、DAT的网络结构设计

2.1 DAT的主要改进

DAT引入了两个重要的创新:可变形注意力机制动态采样点。这两个改进共同作用,使得DAT在处理图像时能够更加集中于有效信息,避免了无效计算。

  • 可变形注意力:通过动态调整采样点的方式,DAT只处理图像中的关键区域,从而减少了无关区域的计算负担。这种方式不仅能够保持良好的模型性能,还大幅降低了计算复杂度。
  • 动态采样点:DAT根据图像内容动态选择注意力的采样点位置,进一步提升了模型的灵活性和适应性。

2.2 DAT网络结构示意图

下图展示了DAT的网络结构及其工作原理:

  1. 可变形注意力机制:在图像特征图上,DAT通过引入一组参考点,并根据查询点通过偏移网络学习得到采样点的偏移量。通过这种方式,DAT能够动态生成采样点并进行特征提取。
  2. 偏移生成网络:偏移生成网络负责计算采样点的偏移量,结合参考点动态调整注意力的焦点。

这种动态调整采样点的方式,使得DAT能够根据输入图像的不同区域,灵活调整计算重点,从而提升目标检测任务的效率和性能。

2.3 DAT与传统注意力机制对比

DAT与传统的Transformer自注意力机制和DCN相比,具有更高的计算效率和灵活性。下图展示了DAT与传统机制的对比,可以直观地看到DAT在处理复杂图像任务时如何通过动态采样提高性能。

通过动态采样点,DAT能够跳过无用的区域,集中资源在更具信息性的部分,从而提升模型的检测精度和推理速度。


三、DAT的代码实现详解

在本节中,我们将详细解析DAT的核心代码实现,并为每一行代码添加注释,帮助您深入理解其工作原理。以下代码展示了DAT的网络结构及其核心机制。

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import einops  # 用于高效的张量操作
from timm.models.layers import to_2tuple, trunc_normal_

# 定义LayerNormProxy类,用于将LayerNorm应用于四维张量
class LayerNormProxy(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)  # 定义LayerNorm

    def forward(self, x):
        x = einops.rearrange(x, 'b c h w -> b h w c')  # 调整张量形状以适应LayerNorm
        x = self.norm(x)  # 应用LayerNorm
        return einops.rearrange(x, 'b h w c -> b c h w')  # 恢复张量形状

# 定义DAT的核心类DAttentionBaseline
class DAttentionBaseline(nn.Module):
    def __init__(self, q_size=(224,224), kv_size=(224,224), n_heads=8, n_head_channels=32, n_groups=1,
                 attn_drop=0.0, proj_drop=0.0, stride=1, offset_range_factor=-1, use_pe=True, dwc_pe=True,
                 no_off=False, fixed_pe=False, ksize=9, log_cpb=False):
        super().__init__()

        # 初始化各类参数
        self.dwc_pe = dwc_pe
        self.n_head_channels = n_head_channels
        self.scale = self.n_head_channels ** -0.5  # 缩放因子
        self.n_heads = n_heads
        self.q_h, self.q_w = q_size  # 查询点的高度和宽度
        self.kv_h, self.kv_w = self.q_h // stride, self.q_w // stride  # 键值点的高度和宽度
        self.nc = n_head_channels * n_heads  # 总通道数
        self.n_groups = n_groups  # 分组数
        self.n_group_channels = self.nc // self.n_groups  # 每组的通道数
        self.n_group_heads = self.n_heads // self.n_groups  # 每组的头数
        self.use_pe = use_pe  # 是否使用位置嵌入
        self.fixed_pe = fixed_pe  # 是否使用固定位置嵌入
        self.no_off = no_off  # 是否禁用偏移
        self.offset_range_factor = offset_range_factor  # 偏移范围系数
        self.ksize = ksize  # 卷积核大小
        self.log_cpb = log_cpb  # 是否使用log-cpb
        self.stride = stride  # 步长
        pad_size = ksize // 2 if ksize != stride else 0  # 计算填充大小

        # 定义偏移卷积层
        self.conv_offset = nn.Sequential(
            nn.Conv2d(self.n_group_channels, self.n_group_channels, ksize, stride, pad_size, groups=self.n_group_channels),
            LayerNormProxy(self.n_group_channels),  # 添加LayerNorm层
            nn.GELU(),  # 使用GELU激活函数
            nn.Conv2d(self.n_group_channels, 2, 1, 1, 0, bias=False) 

 # 输出偏移量
        )

        # 如果禁用偏移,禁用偏移层的参数更新
        if self.no_off:
            for m in self.conv_offset.parameters():
                m.requires_grad_(False)

        # 定义查询、键和值的投影层
        self.proj_q = nn.Conv2d(self.nc, self.nc, kernel_size=1, stride=1, padding=0)
        self.proj_k = nn.Conv2d(self.nc, self.nc, kernel_size=1, stride=1, padding=0)
        self.proj_v = nn.Conv2d(self.nc, self.nc, kernel_size=1, stride=1, padding=0)
        self.proj_out = nn.Conv2d(self.nc, self.nc, kernel_size=1, stride=1, padding=0)

        # 定义丢弃层
        self.proj_drop = nn.Dropout(proj_drop, inplace=True)
        self.attn_drop = nn.Dropout(attn_drop, inplace=True)

        # 定义位置嵌入
        if self.use_pe and not self.no_off:
            if self.dwc_pe:
                self.rpe_table = nn.Conv2d(self.nc, self.nc, kernel_size=3, stride=1, padding=1, groups=self.nc)
            elif self.fixed_pe:
                self.rpe_table = nn.Parameter(torch.zeros(self.n_heads, self.q_h * self.q_w, self.kv_h * self.kv_w))
                trunc_normal_(self.rpe_table, std=0.01)
            elif self.log_cpb:
                self.rpe_table = nn.Sequential(
                    nn.Linear(2, 32, bias=True),
                    nn.ReLU(inplace=True),
                    nn.Linear(32, self.n_group_heads, bias=False)
                )
            else:
                self.rpe_table = nn.Parameter(torch.zeros(self.n_heads, self.q_h * 2 - 1, self.q_w * 2 - 1))
                trunc_normal_(self.rpe_table, std=0.01)
        else:
            self.rpe_table = None

    # 获取参考点
    @torch.no_grad()
    def _get_ref_points(self, H_key, W_key, B, dtype, device):
        ref_y, ref_x = torch.meshgrid(
            torch.linspace(0.5, H_key - 0.5, H_key, dtype=dtype, device=device),
            torch.linspace(0.5, W_key - 0.5, W_key, dtype=dtype, device=device),
            indexing='ij'
        )
        ref = torch.stack((ref_y, ref_x), -1)
        ref[..., 1].div_(W_key - 1.0).mul_(2.0).sub_(1.0)
        ref[..., 0].div_(H_key - 1.0).mul_(2.0).sub_(1.0)
        ref = ref[None, ...].expand(B * self.n_groups, -1, -1, -1)  # 扩展维度
        return ref

    # 获取查询点网格
    @torch.no_grad()
    def _get_q_grid(self, H, W, B, dtype, device):
        ref_y, ref_x = torch.meshgrid(
            torch.arange(0, H, dtype=dtype, device=device),
            torch.arange(0, W, dtype=dtype, device=device),
            indexing='ij'
        )
        ref = torch.stack((ref_y, ref_x), -1)
        ref[..., 1].div_(W - 1.0).mul_(2.0).sub_(1.0)
        ref[..., 0].div_(H - 1.0).mul_(2.0).sub_(1.0)
        ref = ref[None, ...].expand(B * self.n_groups, -1, -1, -1)  # 扩展维度
        return ref

    # 前向传播
    def forward(self, x):
        B, C, H, W = x.size()  # 获取输入张量的维度
        dtype, device = x.dtype, x.device  # 获取数据类型和设备信息

        q = self.proj_q(x)  # 计算查询投影
        q_off = einops.rearrange(q, 'b (g c) h w -> (b g) c h w', g=self.n_groups, c=self.n_group_channels)
        offset = self.conv_offset(q_off).contiguous()  # 计算偏移量
        Hk, Wk = offset.size(2), offset.size(3)  # 获取偏移量的高度和宽度
        n_sample = Hk * Wk  # 计算采样点数量

        # 调整偏移量范围
        if self.offset_range_factor >= 0 and not self.no_off:
            offset_range = torch.tensor([1.0 / (Hk - 1.0), 1.0 / (Wk - 1.0)], device=device).reshape(1, 2, 1, 1)
            offset = offset.tanh().mul(offset_range).mul(self.offset_range_factor)

        offset = einops.rearrange(offset, 'b p h w -> b h w p')  # 调整偏移张量形状
        reference = self._get_ref_points(Hk, Wk, B, dtype, device)  # 获取参考点

        # 如果禁用偏移,使用平均池化替代
        if self.no_off:
            x_sampled = F.avg_pool2d(x, kernel_size=self.stride, stride=self.stride)
            assert x_sampled.size(2) == Hk and x_sampled.size(3) == Wk, f"Size is {x_sampled.size()}"
        else:
            # 使用偏移量和参考点进行特征采样
            x_sampled = F.grid_sample(
                input=x.reshape(B * self.n_groups, self.n_group_channels, H, W),
                grid=pos[..., (1, 0)],  # 使用偏移的坐标进行采样
                mode='bilinear', align_corners=True)  # 双线性插值

        x_sampled = x_sampled.reshape(B, C, 1, n_sample)  # 调整采样后的张量形状

        # 计算键和值的投影
        k = self.proj_k(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample)
        v = self.proj_v(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample)

        # 计算注意力得分
        attn = torch.einsum('b c m, b c n -> b m n', q, k)  # B * h, HW, Ns
        attn = attn.mul(self.scale)  # 缩放注意力得分

        # 如果使用位置嵌入,添加位置偏置
        if self.use_pe and (not self.no_off):
            if self.dwc_pe:
                residual_lepe = self.rpe_table(q.reshape(B, C, H, W)).reshape(B * self.n_heads, self.n_head_channels, H * W)
            elif self.fixed_pe:
                attn_bias = self.rpe_table[None, ...].expand(B, -1, -1, -1)
                attn = attn + attn_bias.reshape(B * self.n_heads, H * W, n_sample)

        attn = F.softmax(attn, dim=2)  # 归一化注意力得分
        attn = self.attn_drop(attn)  # 应用丢弃层

        # 计算注意力输出
        out = torch.einsum('b m n, b c n -> b c m', attn, v)
        out = out.reshape(B, C, H, W)

        # 如果使用深度可分离位置嵌入,添加残差
        if self.use_pe and self.dwc_pe:
            out = out + residual_lepe
        y = self.proj_drop(self.proj_out(out))  # 输出经过投影后的结果
        return y  # 返回最终输出

3.1 代码详解

  1. LayerNormProxy:该类用于在四维张量上应用LayerNorm操作,它通过调整输入张量的形状使其适应LayerNorm层的计算。
  2. DAttentionBaseline:这是DAT的核心类,定义了可变形注意力的各个部分,包括查询、键、值的投影,偏移卷积,以及最终的特征采样和注意力计算。通过动态调整偏移量和采样点,DAT能够高效地处理不同的图像区域。
  3. 位置嵌入:DAT可以选择性地使用位置嵌入(PE),以增强模型的空间感知能力。在这里,我们提供了多种位置嵌入方式,如深度可分离卷积、固定PE和log-cpb等。

四、将DAT集成到YOLOv8中

将DAT集成到YOLOv8的步骤并不复杂,您可以按照以下几个步骤进行操作:

4.1 修改模型配置

首先,您需要修改YOLOv8的模型配置文件,找到对应的yaml文件,例如yolov8.yaml。在该文件中,

找到网络的backboneneck部分,并将DAT模块添加到适当的位置。例如,您可以在主干网络中的某一层加入DAT模块,或者在Neck部分的SPPF模块前后加入。

4.2 修改损失函数代码

在YOLOv8的损失函数文件中(通常位于loss.py),找到与模型输出相关的代码,并确保模型输出在加入DAT之后能够正确传递到损失函数中。您可能需要根据DAT模块的输出形式调整相应的代码逻辑。

4.3 训练和评估模型

完成DAT的集成之后,您可以按照正常的流程训练和评估模型。建议在小规模数据集上进行初步测试,确保DAT的工作正常。之后,您可以在更大规模的数据集上进行完整训练,观察性能提升情况。


五、DAT可添加的位置及推荐策略

5.1 DAT的推荐添加位置

DAT作为一种灵活的注意力机制,可以在模型的多个位置进行集成。以下是一些推荐的添加位置:

  1. 残差连接中:在YOLOv8的残差模块中加入DAT,可以有效提升模型的特征提取能力。
  2. 特征金字塔(SPPF)之前:在特征金字塔模块之前加入DAT,可以帮助模型更好地融合不同尺度的特征信息。
  3. Neck部分:在Neck部分加入DAT,有助于提高模型在多尺度下的特征融合能力。
  4. 输出层之前:在模型的输出层之前加入DAT,可以使模型在做出最终预测时,更加专注于关键信息。

5.2 DAT添加位置的可视化

下图展示了YOLOv8网络中可能添加DAT的几个关键位置。通过将DAT集成到这些位置,您可以根据不同的任务需求进行灵活调整,以达到最优的模型性能。


六、总结

本文详细介绍了DAT(Deformable Attention Transformer)在YOLOv8中的应用,涵盖了DAT的核心思想、网络结构、代码实现以及如何将其集成到YOLOv8中。通过引入DAT,YOLOv8在处理复杂场景和小目标检测任务时,可以更加高效、精准地捕捉图像中的关键信息。

如果您正在寻求进一步提升YOLOv8模型的检测精度,特别是在处理多尺度或复杂场景时,DAT无疑是一种极具潜力的优化策略。希望通过本文的讲解,您能够顺利实现DAT的集成,并在实际项目中取得更好的检测效果。


参考文献

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐