rpotter6298
127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
"""se_block — Squeeze-and-Excitation channel gating."""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class SEBlock(nn.Module):
|
|
"""SE-style channel gating for vectors, feature maps, and sequences.
|
|
|
|
Input shapes:
|
|
[N, C] (vector) — squeeze is identity
|
|
[N, C, H, W] (image map) — squeeze over H, W
|
|
[N, T, C] (sequence) — squeeze over T
|
|
|
|
Gate modes:
|
|
residual (default): gate = 1 + tanh(MLP(s)) in (0, 2), identity at init
|
|
plain: gate = sigmoid(MLP(s)) in (0, 1)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
reduction: int = 16,
|
|
residual: bool = True,
|
|
identity_init: bool = True,
|
|
):
|
|
super().__init__()
|
|
hid = max(1, dim // max(1, reduction))
|
|
self.fc1 = nn.Linear(dim, hid, bias=True)
|
|
self.act = nn.ReLU(inplace=True)
|
|
self.fc2 = nn.Linear(hid, dim, bias=True)
|
|
self.residual = residual
|
|
|
|
if residual and identity_init:
|
|
nn.init.zeros_(self.fc2.weight)
|
|
nn.init.zeros_(self.fc2.bias)
|
|
|
|
def _squeeze(self, x: torch.Tensor) -> torch.Tensor:
|
|
if x.dim() == 2:
|
|
return x
|
|
if x.dim() == 4:
|
|
return x.mean(dim=(2, 3))
|
|
if x.dim() == 3:
|
|
return x.mean(dim=1)
|
|
return x.view(x.size(0), -1)
|
|
|
|
def _broadcast(self, gate: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
|
|
if like.dim() == 2:
|
|
return gate
|
|
if like.dim() == 3:
|
|
return gate.unsqueeze(1)
|
|
if like.dim() == 4:
|
|
return gate.unsqueeze(-1).unsqueeze(-1)
|
|
return gate.view_as(like)
|
|
|
|
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
s = self._squeeze(x)
|
|
u = self.fc2(self.act(self.fc1(s)))
|
|
gate = (1.0 + torch.tanh(u)) if self.residual else torch.sigmoid(u)
|
|
return x * self._broadcast(gate, x), gate
|
|
|
|
|
|
class SEGateLogger:
|
|
"""Running stats over SE gate activations across batches."""
|
|
|
|
def __init__(
|
|
self,
|
|
enabled: bool = True,
|
|
track_channels: bool = False,
|
|
dim: int | None = None,
|
|
):
|
|
self.enabled = enabled
|
|
self.track_channels = track_channels
|
|
self.dim = dim
|
|
self.reset()
|
|
|
|
def reset(self) -> None:
|
|
self._n = 0
|
|
self._sum = 0.0
|
|
self._sum2 = 0.0
|
|
self._lt02 = 0
|
|
self._gt08 = 0
|
|
self._ch_sum = None
|
|
self._ch_count = 0
|
|
if self.track_channels and self.dim is not None:
|
|
self._ch_sum = torch.zeros(self.dim, dtype=torch.float32)
|
|
|
|
@torch.no_grad()
|
|
def accumulate(self, gates: torch.Tensor) -> None:
|
|
if not self.enabled:
|
|
return
|
|
if gates.dim() == 4:
|
|
g = gates.mean(dim=(2, 3))
|
|
elif gates.dim() == 3:
|
|
g = gates.mean(dim=1)
|
|
elif gates.dim() == 2:
|
|
g = gates
|
|
else:
|
|
g = gates.view(gates.size(0), -1)
|
|
g = g.detach()
|
|
self._n += g.numel()
|
|
self._sum += g.sum().item()
|
|
self._sum2 += (g * g).sum().item()
|
|
self._lt02 += (g < 0.2).sum().item()
|
|
self._gt08 += (g > 0.8).sum().item()
|
|
if self._ch_sum is not None:
|
|
self._ch_sum += g.sum(dim=0).cpu()
|
|
self._ch_count += g.size(0)
|
|
|
|
def get(self, reset: bool = True) -> dict | None:
|
|
if self._n == 0:
|
|
return None
|
|
mean = self._sum / self._n
|
|
var = max(0.0, self._sum2 / self._n - mean * mean)
|
|
out = {
|
|
"mean": mean,
|
|
"std": var ** 0.5,
|
|
"pct_lt_0.2": self._lt02 / self._n,
|
|
"pct_gt_0.8": self._gt08 / self._n,
|
|
}
|
|
if self._ch_sum is not None and self._ch_count > 0:
|
|
out["channel_mean"] = (self._ch_sum / float(self._ch_count)).tolist()
|
|
if reset:
|
|
self.reset()
|
|
return out
|