File size: 1,380 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 | """ECA (Efficient Channel Attention) — 残差初始化版本，注册到 ultralytics
关键改进：使用可学习缩放因子 alpha (初始=0)，使得
output = x * (1 + alpha * (sigmoid(y) - 1))
当 alpha=0 时 output=x（恒等映射），不破坏预训练特征。
"""
import math
import torch
import torch.nn as nn
class ECA(nn.Module):
def __init__(self, c1=None, c2=None, gamma=2, b=1):
super().__init__()
self.gamma = gamma
self.b = b
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.sigmoid = nn.Sigmoid()
self.alpha = nn.Parameter(torch.zeros(1))
self._conv = None
if c1 is not None:
self._init_conv(c1)
def _init_conv(self, channels):
t = int(abs((math.log2(channels) + self.b) / self.gamma))
k = t if t % 2 else t + 1
self._conv = nn.Conv1d(1, 1, kernel_size=k, padding=k // 2, bias=False)
def forward(self, x):
if self._conv is None:
self._init_conv(x.shape[1])
self._conv = self._conv.to(x.device)
y = self.avg_pool(x)
y = self._conv(y.squeeze(-1).transpose(-1, -2))
y = y.transpose(-1, -2).unsqueeze(-1)
return x * (1.0 + self.alpha * (self.sigmoid(y) - 1.0))
def register_eca():
import ultralytics.nn.tasks as tasks
if "ECA" not in vars(tasks):
tasks.ECA = ECA
|