File size: 2,285 Bytes
6a5bb7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """CBAM (Convolutional Block Attention Module) — 残差初始化版本
通道注意力：学习哪些特征通道更重要（类似 SE/ECA）
空间注意力：学习哪些空间位置更重要（对遮挡、密集场景有帮助）
残差初始化：alpha 初始=0 → 输出恒等，不破坏预训练特征。
"""
import torch
import torch.nn as nn
class ChannelAttention(nn.Module):
def __init__(self, channels, reduction=16):
super().__init__()
mid = max(channels // reduction, 8)
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.max_pool = nn.AdaptiveMaxPool2d(1)
self.fc = nn.Sequential(
nn.Conv2d(channels, mid, 1, bias=False),
nn.ReLU(inplace=True),
nn.Conv2d(mid, channels, 1, bias=False),
)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
return self.sigmoid(self.fc(self.avg_pool(x)) + self.fc(self.max_pool(x)))
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super().__init__()
self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
return self.sigmoid(self.conv(torch.cat([avg_out, max_out], dim=1)))
class CBAM(nn.Module):
def __init__(self, c1=None, c2=None, reduction=16, kernel_size=7):
super().__init__()
self.reduction = reduction
self.kernel_size = kernel_size
self.alpha = nn.Parameter(torch.zeros(1))
self.ca = None
self.sa = None
if c1 is not None:
self._init_modules(c1)
def _init_modules(self, channels):
self.ca = ChannelAttention(channels, self.reduction)
self.sa = SpatialAttention(self.kernel_size)
def forward(self, x):
if self.ca is None:
self._init_modules(x.shape[1])
self.ca = self.ca.to(x.device)
self.sa = self.sa.to(x.device)
attn = self.ca(x) * self.sa(x)
return x * (1.0 + self.alpha * (attn - 1.0))
def register_cbam():
import ultralytics.nn.tasks as tasks
if "CBAM" not in vars(tasks):
tasks.CBAM = CBAM
|