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