1.PolarizedSelfAttention
Polarized self-attention是一种改进的自注意力机制,旨在提高自注意力模型的效率和准确性。传统的自注意力机制通常需要计算所有位置对所有位置的注意力,这会导致计算复杂度的增加和模型的训练时间的延长。Polarized self-attention通过引入极化因子来减少需要计算的位置对之间的注意力,从而提高自注意力模型的效率。极化因子是一个向量,可以将其与每个位置的向量点积,以确定哪些位置需要进行注意力计算。这种方法可以通过在保持准确度的同时减少计算量来提高自注意力模型的效率。
基于双重注意力机制,本文针对Pixel-wise regression的任务,提出了一种更加精细的双重注意力机制——极化自注意力(Polarized Self-Attention)。作为一个即插即用的模块,在人体姿态估计和语义分割任务上,作者将它用在了以前的SOTA模型上,并达到了新的SOTA性能,霸榜COCO人体姿态估计和Cityscapes语义分割。
注意力机制是一个被广泛应用在各种CV任务中的方法。注意力机制根据施加的维度大致可以分为两类:通道注意力和空间注意力。对于通道注意力机制,代表性的工作有SENet、ECANet;对于空间注意力机制,代表性的工作有Self-Attention。随着空间和通道注意力机制的提出,很自然的,结合空间和通道两个维度的双重注意力机制也被提出,代表工作有CBAM,DANet。
作者提出了一个即插即用的模块——极化自注意力机制( Polarized Self-Attention(PSA)),用于解决像素级的回归任务,相比于其他注意力机制,极化自注意力机制主要有两个设计上的亮点:
1)极化滤波( Polarized filtering):在通道和空间维度保持比较高的resolution(在通道上保持C/2的维度,在空间上保持[H,W]的维度 ),这一步能够减少降维度造成的信息损失;
2)增强(Enhancement):组合非线性直接拟合典型细粒度回归的输出分布。
###################### PolarizedSelfAttention #### start ###############################
from torch import nn
from torch.nn import init
class PolarizedSelfAttention(nn.Module):
def __init__(self, channel=512):
super().__init__()
self.ch_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))
self.ch_wq = nn.Conv2d(channel, 1, kernel_size=(1, 1))
self.softmax_channel = nn.Softmax(1)
self.softmax_spatial = nn.Softmax(-1)
self.ch_wz = nn.Conv2d(channel // 2, channel, kernel_size=(1, 1))
self.ln = nn.LayerNorm(channel)
self.sigmoid = nn.Sigmoid()
self.sp_wv = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))
self.sp_wq = nn.Conv2d(channel, channel // 2, kernel_size=(1, 1))
self.agp = nn.AdaptiveAvgPool2d((1, 1))
def forward(self, x):
b, c, h, w = x.size()
# Channel-only Self-Attention
channel_wv = self.ch_wv(x) # bs,c//2,h,w
channel_wq = self.ch_wq(x) # bs,1,h,w
channel_wv = channel_wv.reshape(b, c // 2, -1) # bs,c//2,h*w
channel_wq = channel_wq.reshape(b, -1, 1) # bs,h*w,1
channel_wq = self.softmax_channel(channel_wq)
channel_wz = torch.matmul(channel_wv, channel_wq).unsqueeze(-1) # bs,c//2,1,1
channel_weight = self.sigmoid(self.ln(self.ch_wz(channel_wz).reshape(b, c, 1).permute(0, 2, 1))).permute(0, 2,
1).reshape(
b, c, 1, 1) # bs,c,1,1
channel_out = channel_weight * x
# Spatial-only Self-Attention
spatial_wv = self.sp_wv(x) # bs,c//2,h,w
spatial_wq = self.sp_wq(x) # bs,c//2,h,w
spatial_wq = self.agp(spatial_wq) # bs,c//2,1,1
spatial_wv = spatial_wv.reshape(b, c // 2, -1) # bs,c//2,h*w
spatial_wq = spatial_wq.permute(0, 2, 3, 1).reshape(b, 1, c // 2) # bs,1,c//2
spatial_wq = self.softmax_spatial(spatial_wq)
spatial_wz = torch.matmul(spatial_wq, spatial_wv) # bs,1,h*w
spatial_weight = self.sigmoid(spatial_wz.reshape(b, 1, h, w)) # bs,1,h,w
spatial_out = spatial_weight * x
out = spatial_out + channel_out
return out
###################### PolarizedSelfAttention #### end ###############################
首先PolarizedSelfAttention进行注册
from ultralytics.nn.attention.attention import *
函数def parse_model(d, ch, verbose=True): # model_dict, input_channels(3)进行修改
#####attention ####
elif m in (MHSA,ECAAttention,TripletAttention,BAM,PolarizedSelfAttention):
c1, c2 = ch[f], args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels) * width, 8)
args = [c1, *args[1:]]
#####attention ####
加入backbone SPPF后
# 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 GFLOPs
# 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
- [-1, 1, PolarizedSelfAttention, [1024]] # 10
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 13
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 16 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 13], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 19 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 10], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 22 (P5/32-large)
- [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)
neck里的连接Detect的3个C2f结合
# 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 GFLOPs
# 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, PolarizedSelfAttention, [256]] # 16
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 19 (P4/16-medium)
- [-1, 1, PolarizedSelfAttention, [512]] # 20
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 23 (P5/32-large)
- [-1, 1, PolarizedSelfAttention, [1024]] # 24
- [[16, 20, 24], 1, Detect, [nc]] # Detect(P3, P4, P5)
放入neck的C2f后面
# Ultralytics YOLO 🚀, GPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
# Parameters
nc: 1 # 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 GFLOPs
# 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, PolarizedSelfAttention, [512]] # 13
- [-1, 1, nn.Upsample, [None, 2, 'nearest']]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 16 (P3/8-small)
- [-1, 1, PolarizedSelfAttention, [256]] # 17 (P5/32-large)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 13], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 20 (P4/16-medium)
- [-1, 1, PolarizedSelfAttention, [512]] # 21 (P5/32-large)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 24 (P5/32-large)
- [-1, 1, PolarizedSelfAttention, [1024]] # 25 (P5/32-large)
- [[17, 21, 25], 1, Detect, [nc]] # Detect(P3, P4, P5)
阅读量:268
点赞量:0
收藏量:0