reworked iop_corr, added explainability tools and plotting tools, cleanup codebase

classes/v2/SE_attention.py

@@ -0,0 +1,123 @@
+# se_block.py
+import torch
+import torch.nn as nn
+
+class SEGateLogger:
+    """
+    Lightweight stats over SE gates.
+    Use: logger.accumulate(gates) each batch; logger.get() at epoch end.
+    """
+    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):
+        self._n = 0
+        self._sum = 0.0
+        self._sum2 = 0.0
+        self._lt02 = 0
+        self._gt08 = 0
+        # optional per-channel
+        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):
+        if not self.enabled:
+            return
+        # gates expected shape [N, C]; if a map/sequence gate is passed, reduce to [N, C]
+        if gates.dim() == 4:   # [N,C,H,W] gates (uncommon)
+            g = gates.mean(dim=(2,3))
+        elif gates.dim() == 3: # [N,T,C] gates (sequence)
+            g = gates.mean(dim=1)
+        elif gates.dim() == 2: # [N,C]
+            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):
+        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
+
+class SEBlock(nn.Module):
+    """
+    SE-style channel gating that works for vectors and maps.
+
+    Input:
+      - [N, C]         (vector)      -> squeeze = 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:
+            # make MLP output ~0 at start → gate ≈ 1.0
+            nn.init.zeros_(self.fc2.weight)
+            nn.init.zeros_(self.fc2.bias)
+
+    def _squeeze(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 2:      # [N,C]
+            return x
+        if x.dim() == 4:      # [N,C,H,W]
+            return x.mean(dim=(2,3))
+        if x.dim() == 3:      # [N,T,C]
+            return x.mean(dim=1)
+        # fallback: flatten non-batch dims into channels
+        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)            # [N,1,C]
+        if like.dim() == 4:
+            return gate.unsqueeze(-1).unsqueeze(-1)  # [N,C,1,1]
+        return gate.view_as(like)
+
+    def forward(self, x: torch.Tensor):
+        s = self._squeeze(x)                   # [N,C]
+        u = self.fc2(self.act(self.fc1(s)))    # [N,C]
+        if self.residual:
+            gate = 1.0 + torch.tanh(u)         # (0, 2) with identity at 1.0
+        else:
+            gate = torch.sigmoid(u)            # (0, 1)
+        y = x * self._broadcast(gate, x)
+        return y, gate  # return both the reweighted tensor and the gate for logging