goat / Scripts /modules /eca.py
LightChuan's picture
Upload folder using huggingface_hub
6a5bb7e verified
Raw
History Blame Contribute Delete
1.38 kB
"""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