| """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 |
|
|
