1,本文介绍
Mamba-Like Linear Attention (MLLA)是一种模型,结合了Mamba和线性注意力Transformer的优点,旨在提高模型在视觉任务中的性能。
-
线性注意力:
- 线性注意力是一种注意力机制,用于计算输入序列中各个位置之间的关联性。与传统的Softmax注意力不同,线性注意力使用线性归一化代替非线性Softmax函数,从而降低计算复杂度。这使得线性注意力的计算复杂度从O(N^2)降低到O(N),提高了效率。
-
Mamba模型:
- Mamba是一种基于状态空间的模型,具有线性计算复杂度。它通过有效地建模序列来处理长序列,适用于自然语言处理和视觉识别任务。
-
MLLA的数学原理:
- MLLA模型将Mamba的关键设计元素(如遗忘门和块设计)与线性注意力Transformer相结合。在数学上,这意味着将Mamba的状态空间建模方法与线性注意力的位置关联性计算相结合,以提高模型的表现。
-
遗忘门的替代:
- 在MLLA中,遗忘门通常被替代为适当的位置编码。这些位置编码可以在视觉任务中代替遗忘门的功能,同时保持模型的并行计算和快速推理速度。
-
核心设计元素的整合:
- MLLA模型的关键设计元素包括遗忘门和块设计。通过将这些设计元素整合到线性注力中,MLLA模型能够在图像分类和密集预测任务中取得优越性
关于MLLA的详细介绍可以看论文:https://arxiv.org/pdf/2405.16605
本文将讲解如何将MLLA融合进yolov8
话不多说,上代码!
2 将MLLA融合进yolov8
2.1 步骤一
首先找到如下的目录'ultralytics/nn',然后在这个目录下创建一个'Addmodules'文件夹,然后在这个目录下创建一个MLLA.py文件,文件名字可以根据你自己的习惯起,然后将MLLA的核心代码复制进去。
# MLLA的核心代码
# --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu
# --------------------------------------------------------
# Demystify Mamba in Vision: A Linear Attention Perspective
# Modified by Dongchen Han
# -----------------------------------------------------------------------
import torch
import torch.nn as nn
__all__ = ['MLLAttention']
class Mlp(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
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class ConvLayer(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0, dilation=1, groups=1,
bias=True, dropout=0, norm=nn.BatchNorm2d, act_func=nn.ReLU):
super(ConvLayer, self).__init__()
self.dropout = nn.Dropout2d(dropout, inplace=False) if dropout > 0 else None
self.conv = nn.Conv2d(
in_channels,
out_channels,
kernel_size=(kernel_size, kernel_size),
stride=(stride, stride),
padding=(padding, padding),
dilation=(dilation, dilation),
groups=groups,
bias=bias,
)
self.norm = norm(num_features=out_channels) if norm else None
self.act = act_func() if act_func else None
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.dropout is not None:
x = self.dropout(x)
x = self.conv(x)
if self.norm:
x = self.norm(x)
if self.act:
x = self.act(x)
return x
class RoPE(torch.nn.Module):
r"""Rotary Positional Embedding.
"""
def __init__(self, base=10000):
super(RoPE, self).__init__()
self.base = base
def generate_rotations(self, x):
# 获取输入张量的形状
*channel_dims, feature_dim = x.shape[1:-1][0], x.shape[-1]
k_max = feature_dim // (2 * len(channel_dims))
assert feature_dim % k_max == 0, "Feature dimension must be divisible by 2 * k_max"
# 生成角度
theta_ks = 1 / (self.base ** (torch.arange(k_max, dtype=x.dtype, device=x.device) / k_max))
angles = torch.cat([t.unsqueeze(-1) * theta_ks for t in
torch.meshgrid([torch.arange(d, dtype=x.dtype, device=x.device) for d in channel_dims],
indexing='ij')], dim=-1)
# 计算旋转矩阵的实部和虚部
rotations_re = torch.cos(angles).unsqueeze(dim=-1)
rotations_im = torch.sin(angles).unsqueeze(dim=-1)
rotations = torch.cat([rotations_re, rotations_im], dim=-1)
return rotations
def forward(self, x):
# 生成旋转矩阵
rotations = self.generate_rotations(x)
# 将 x 转换为复数形式
x_complex = torch.view_as_complex(x.reshape(*x.shape[:-1], -1, 2))
# 应用旋转矩阵
pe_x = torch.view_as_complex(rotations) * x_complex
# 将结果转换回实数形式并展平最后两个维度
return torch.view_as_real(pe_x).flatten(-2)
class MLLAttention(nn.Module):
r""" Linear Attention with LePE and RoPE.
Args:
dim (int): Number of input channels.
num_heads (int): Number of attention heads.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
"""
def __init__(self, dim=3, input_resolution=[160, 160], num_heads=4, qkv_bias=True, **kwargs):
super().__init__()
self.dim = dim
self.input_resolution = input_resolution
self.num_heads = num_heads
self.qk = nn.Linear(dim, dim * 2, bias=qkv_bias)
self.elu = nn.ELU()
self.lepe = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
self.rope = RoPE()
def forward(self, x):
"""
Args:
x: input features with shape of (B, N, C)
"""
x = x.reshape((x.size(0), x.size(2) * x.size(3), x.size(1)))
b, n, c = x.shape
h = int(n ** 0.5)
w = int(n ** 0.5)
# self.rope = RoPE(shape=(h, w, self.dim))
num_heads = self.num_heads
head_dim = c // num_heads
qk = self.qk(x).reshape(b, n, 2, c).permute(2, 0, 1, 3)
q, k, v = qk[0], qk[1], x
# q, k, v: b, n, c
q = self.elu(q) + 1.0
k = self.elu(k) + 1.0
q_rope = self.rope(q.reshape(b, h, w, c)).reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)
k_rope = self.rope(k.reshape(b, h, w, c)).reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)
q = q.reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)
k = k.reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)
v = v.reshape(b, n, num_heads, head_dim).permute(0, 2, 1, 3)
z = 1 / (q @ k.mean(dim=-2, keepdim=True).transpose(-2, -1) + 1e-6)
kv = (k_rope.transpose(-2, -1) * (n ** -0.5)) @ (v * (n ** -0.5))
x = q_rope @ kv * z
x = x.transpose(1, 2).reshape(b, n, c)
v = v.transpose(1, 2).reshape(b, h, w, c).permute(0, 3, 1, 2)
x = x + self.lepe(v).permute(0, 2, 3, 1).reshape(b, n, c)
x = x.transpose(2, 1).reshape((b, c, h, w))
return x
def extra_repr(self) -> str:
return f'dim={self.dim}, num_heads={self.num_heads}'
if __name__ == "__main__":
# Generating Sample image
image_size = (1, 64, 160, 160)
image = torch.rand(*image_size)
# Model
model = MLLAttention(64)
out = model(image)
print(out.size())
第二步我们在该目录(Addmodules)下创建一个新的py文件名字为'__init__.py',然后在其内部添加如下代码。
最终结果如下图标注所示
2.2 步骤二
在task.py中进行导入,如下图所示
2.3 步骤三
在task.py的parse_model函数中添加如下代码
2.4 步骤四
找到ultralytics/models/yolo/detect/train.py的DetectionTrainer class中的build_dataset函数中的rect=mode == 'val'改为rect=False
到此注册成功,复制后面的yaml文件直接运行即可
yaml文件
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs
s: [0.33, 0.50, 1024] # YOLOv8s summary: 225 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPs
m: [0.67, 0.75, 768] # YOLOv8m summary: 295 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPs
l: [1.00, 1.00, 512] # YOLOv8l summary: 365 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPs
x: [1.00, 1.25, 512] # YOLOv8x summary: 365 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOP
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [-1, 1, MLLAttention, []] # 22 (P5/32-large) # 添加在大目标检测层后!
- [[15, 18, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)
# 关于MLLA添加的位置还可以放在每个检测头的前面,也可以放在骨干网络的后面
不知不觉已经看完了哦,动动小手留个点赞吧--_--
有问题可以留在评论区哦,博主看到了就会回复