399 lines
14 KiB
Python
399 lines
14 KiB
Python
"""SOFIA v7.1 custom layers.
|
||
|
||
Includes:
|
||
- GGeM: Generalized Mean Pooling with per-channel learnable exponent (F11)
|
||
- CircularHarmonicPool: Formally SO(2)-invariant pooling via polar + FFT magnitude (NH2 novel)
|
||
- AltitudeFiLM: FiLM conditioning on UAV altitude (NH4 novel)
|
||
- RoPE2D: 2D Rotary Position Embedding for attention
|
||
- SqueezeExcite: standard SE block
|
||
- LayerNorm2d: channel-last LN wrapper for 2D features
|
||
|
||
All rotation-invariance and FiLM modules are NOVEL contributions of SOFIA v7.1-α.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from typing import Optional, Tuple
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
|
||
# ============================================================
|
||
# GGeM: Generalized Mean Pooling (F11)
|
||
# ============================================================
|
||
|
||
class GGeM(nn.Module):
|
||
"""Per-channel learnable Generalized Mean pooling.
|
||
|
||
Formula:
|
||
GGeM(F)_c = (1/HW · Σ F_{c,h,w}^{p_c})^{1/p_c}
|
||
p_c = softplus(p_hat_c) ∈ (0, ∞)
|
||
"""
|
||
|
||
def __init__(self, channels: int, init_p: float = 3.0, eps: float = 1e-6) -> None:
|
||
super().__init__()
|
||
# softplus^{-1}(init_p) = log(exp(init_p) - 1)
|
||
hat_init = math.log(math.exp(init_p) - 1.0)
|
||
self.hat_p = nn.Parameter(torch.full((channels,), hat_init))
|
||
self.eps = eps
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""x: [B, C, H, W] -> [B, C]"""
|
||
p = F.softplus(self.hat_p).view(1, -1, 1, 1) # [1, C, 1, 1]
|
||
x_clamped = x.clamp(min=self.eps)
|
||
x_pow = x_clamped.pow(p)
|
||
x_mean = x_pow.mean(dim=(2, 3), keepdim=True)
|
||
out = x_mean.pow(1.0 / p)
|
||
return out.flatten(1)
|
||
|
||
|
||
# ============================================================
|
||
# CHP: Circular Harmonic Pool (NH2 novel — formally SO(2)-invariant)
|
||
# ============================================================
|
||
|
||
class CircularHarmonicPool(nn.Module):
|
||
"""Formally SO(2) rotation-invariant pooling.
|
||
|
||
Algorithm:
|
||
1. Sample input feature map at polar grid (r, θ) via bilinear grid_sample
|
||
2. Apply 1D real FFT along θ-axis
|
||
3. Keep magnitudes of first N harmonics (invariant to shift = rotation)
|
||
4. GGeM pool over rings r (per-channel-per-harmonic)
|
||
5. Flatten to descriptor [B, C * N]
|
||
|
||
Output is theoretically invariant to input rotation of any angle.
|
||
|
||
See HYP Phase 4'' Section 4''.1 NH2 and Section 4''.5 for formal proof.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
channels: int,
|
||
rings: int = 8,
|
||
angles: int = 16,
|
||
harmonics: int = 4,
|
||
r_min: float = 0.1,
|
||
r_max: float = 1.0,
|
||
) -> None:
|
||
super().__init__()
|
||
assert harmonics <= angles // 2 + 1, (
|
||
f"harmonics {harmonics} cannot exceed angles//2+1 = {angles // 2 + 1}"
|
||
)
|
||
|
||
self.channels = channels
|
||
self.rings = rings
|
||
self.angles = angles
|
||
self.harmonics = harmonics
|
||
|
||
# GGeM over rings (per channel × per harmonic)
|
||
self.ggem = GGeM(channels * harmonics, init_p=3.0)
|
||
|
||
# Precompute polar grid in normalized [-1, 1] coords for grid_sample
|
||
grid = self._make_polar_grid(rings, angles, r_min, r_max) # [R, T, 2]
|
||
self.register_buffer("polar_grid", grid, persistent=False)
|
||
|
||
@staticmethod
|
||
def _make_polar_grid(R: int, T: int, r_min: float, r_max: float) -> torch.Tensor:
|
||
r_values = torch.linspace(r_min, r_max, R) # [R]
|
||
theta_values = torch.linspace(0.0, 2 * math.pi, T + 1)[:-1] # [T]
|
||
r_grid, theta_grid = torch.meshgrid(r_values, theta_values, indexing="ij") # [R, T]
|
||
x = r_grid * torch.cos(theta_grid)
|
||
y = r_grid * torch.sin(theta_grid)
|
||
grid = torch.stack([x, y], dim=-1) # [R, T, 2]
|
||
return grid
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
Args:
|
||
x: [B, C, H, W] feature map
|
||
Returns:
|
||
descriptor: [B, C * harmonics]
|
||
"""
|
||
B, C, H, W = x.shape
|
||
assert C == self.channels, (
|
||
f"Input channels {C} != expected {self.channels}"
|
||
)
|
||
|
||
# 1. Polar sampling
|
||
# grid: [B, R, T, 2]
|
||
grid = self.polar_grid.unsqueeze(0).expand(B, -1, -1, -1)
|
||
polar = F.grid_sample(
|
||
x, grid,
|
||
mode="bilinear",
|
||
padding_mode="zeros",
|
||
align_corners=True,
|
||
) # [B, C, R, T]
|
||
|
||
# 2. 1D real FFT along angular axis
|
||
polar_fft = torch.fft.rfft(polar, dim=-1) # [B, C, R, T//2+1]
|
||
polar_fft = polar_fft[..., : self.harmonics] # [B, C, R, N]
|
||
|
||
# 3. Magnitude (rotation invariant)
|
||
magnitude = polar_fft.abs() # [B, C, R, N]
|
||
|
||
# 4. Reshape for GGeM: treat (C, N) as combined channel dim, rings as spatial
|
||
# Shape: [B, C*N, R, 1]
|
||
magnitude_reshaped = (
|
||
magnitude
|
||
.permute(0, 1, 3, 2) # [B, C, N, R]
|
||
.reshape(B, C * self.harmonics, self.rings, 1)
|
||
)
|
||
|
||
# 5. GGeM pool over rings (H=R, W=1)
|
||
descriptor = self.ggem(magnitude_reshaped) # [B, C*N]
|
||
|
||
return descriptor
|
||
|
||
|
||
# ============================================================
|
||
# FiLM: Altitude-conditioned modulation (NH4 novel)
|
||
# ============================================================
|
||
|
||
class AltitudeFiLM(nn.Module):
|
||
"""FiLM modulation conditioned on scalar altitude.
|
||
|
||
F' = γ(h) · F + β(h) where γ,β ∈ R^C are produced by MLP.
|
||
|
||
At altitude=None, produces identity (γ=1, β=0) via zero-init of final layer.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
channels: int,
|
||
hidden_dim: int = 64,
|
||
altitude_norm: float = 500.0,
|
||
) -> None:
|
||
super().__init__()
|
||
self.channels = channels
|
||
self.altitude_norm = altitude_norm
|
||
|
||
self.mlp = nn.Sequential(
|
||
nn.Linear(1, hidden_dim),
|
||
nn.GELU(),
|
||
nn.Linear(hidden_dim, 2 * channels),
|
||
)
|
||
# Zero-init last layer → initial γ=0 before residual, β=0
|
||
nn.init.zeros_(self.mlp[-1].weight)
|
||
nn.init.zeros_(self.mlp[-1].bias)
|
||
|
||
def forward(self, x: torch.Tensor, altitude: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||
"""
|
||
Args:
|
||
x: [B, C, H, W]
|
||
altitude: [B] or [B, 1] scalar altitude in meters (or None for neutral)
|
||
Returns:
|
||
[B, C, H, W]
|
||
"""
|
||
B = x.shape[0]
|
||
if altitude is None:
|
||
altitude = torch.zeros(B, 1, device=x.device, dtype=x.dtype)
|
||
elif altitude.dim() == 1:
|
||
altitude = altitude.unsqueeze(-1)
|
||
|
||
h_norm = altitude.to(x.dtype) / self.altitude_norm
|
||
gamma_beta = self.mlp(h_norm) # [B, 2C]
|
||
gamma, beta = gamma_beta.chunk(2, dim=-1)
|
||
# Residual form: γ = 1 + delta_γ (starts at identity)
|
||
gamma = gamma.view(B, self.channels, 1, 1) + 1.0
|
||
beta = beta.view(B, self.channels, 1, 1)
|
||
return gamma * x + beta
|
||
|
||
|
||
# ============================================================
|
||
# TextFiLM: text-conditioned modulation (extension for caption fusion)
|
||
# ============================================================
|
||
|
||
class TextFiLM(nn.Module):
|
||
"""FiLM modulation conditioned on a text embedding.
|
||
|
||
F' = γ(z) · F + β(z) where γ,β ∈ R^C are produced by an MLP
|
||
from z ∈ R^{D_txt}. Identity at init via zero-init of last layer
|
||
(so γ=1, β=0 before training shifts the residual).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
channels: int,
|
||
text_dim: int = 1024,
|
||
hidden_dim: int = 256,
|
||
) -> None:
|
||
super().__init__()
|
||
self.channels = channels
|
||
self.text_dim = text_dim
|
||
self.mlp = nn.Sequential(
|
||
nn.Linear(text_dim, hidden_dim),
|
||
nn.GELU(),
|
||
nn.Linear(hidden_dim, 2 * channels),
|
||
)
|
||
nn.init.zeros_(self.mlp[-1].weight)
|
||
nn.init.zeros_(self.mlp[-1].bias)
|
||
|
||
def forward(self, x: torch.Tensor, text_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||
"""
|
||
Args:
|
||
x: [B, C, H, W]
|
||
text_emb: [B, D_txt] or None for no-op (identity)
|
||
Returns:
|
||
[B, C, H, W]
|
||
"""
|
||
if text_emb is None:
|
||
return x
|
||
B = x.shape[0]
|
||
gamma_beta = self.mlp(text_emb.to(x.dtype)) # [B, 2C]
|
||
gamma, beta = gamma_beta.chunk(2, dim=-1)
|
||
gamma = gamma.view(B, self.channels, 1, 1) + 1.0
|
||
beta = beta.view(B, self.channels, 1, 1)
|
||
return gamma * x + beta
|
||
|
||
|
||
# ============================================================
|
||
# RoPE 2D
|
||
# ============================================================
|
||
|
||
class RoPE2D(nn.Module):
|
||
"""2D Rotary Position Embedding.
|
||
|
||
Splits head_dim into two halves: first half gets x-position encoding,
|
||
second half gets y-position encoding. For each half, applies standard
|
||
1D RoPE rotation.
|
||
|
||
Reference: RoFormer (B49) adapted for 2D.
|
||
"""
|
||
|
||
def __init__(self, head_dim: int, max_resolution: int = 64, base: float = 10000.0) -> None:
|
||
super().__init__()
|
||
assert head_dim % 2 == 0, "head_dim must be even for RoPE"
|
||
assert head_dim % 4 == 0, "head_dim must be divisible by 4 for 2D RoPE"
|
||
self.head_dim = head_dim
|
||
self.half_dim = head_dim // 2 # dedicated to each axis
|
||
self.max_resolution = max_resolution
|
||
|
||
# Frequencies for each axis (half_dim per axis, sin+cos pairs)
|
||
freqs = 1.0 / (base ** (torch.arange(0, self.half_dim, 2).float() / self.half_dim))
|
||
self.register_buffer("freqs", freqs, persistent=False)
|
||
|
||
def _make_embeds(self, H: int, W: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""Produce cos/sin embeddings for HW tokens in raster order."""
|
||
y_pos = torch.arange(H, device=device, dtype=torch.float32)
|
||
x_pos = torch.arange(W, device=device, dtype=torch.float32)
|
||
|
||
freqs_y = torch.einsum("i,j->ij", y_pos, self.freqs) # [H, half_dim/2]
|
||
freqs_x = torch.einsum("i,j->ij", x_pos, self.freqs) # [W, half_dim/2]
|
||
|
||
# Expand to full grid: [H, W, half_dim/2] each
|
||
freqs_y = freqs_y.unsqueeze(1).expand(-1, W, -1) # [H, W, half_dim/2]
|
||
freqs_x = freqs_x.unsqueeze(0).expand(H, -1, -1) # [H, W, half_dim/2]
|
||
|
||
# Concatenate: x-axis freqs into first half, y-axis into second half
|
||
# Each half is [H, W, half_dim/2]; we pair-up for complex rotation
|
||
freqs_combined_x = torch.cat([freqs_x, freqs_x], dim=-1) # [H, W, half_dim]
|
||
freqs_combined_y = torch.cat([freqs_y, freqs_y], dim=-1) # [H, W, half_dim]
|
||
|
||
freqs_full = torch.cat([freqs_combined_x, freqs_combined_y], dim=-1) # [H, W, head_dim]
|
||
cos = freqs_full.cos().reshape(H * W, -1)
|
||
sin = freqs_full.sin().reshape(H * W, -1)
|
||
return cos, sin
|
||
|
||
@staticmethod
|
||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||
"""Rotate: (x1, x2) -> (-x2, x1)."""
|
||
x1, x2 = x.chunk(2, dim=-1)
|
||
return torch.cat([-x2, x1], dim=-1)
|
||
|
||
def forward(self, q: torch.Tensor, k: torch.Tensor, H: int, W: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""
|
||
Args:
|
||
q, k: [B, heads, HW, head_dim]
|
||
H, W: spatial dims for position computation
|
||
Returns:
|
||
q_rot, k_rot with positional encoding applied
|
||
"""
|
||
cos, sin = self._make_embeds(H, W, q.device)
|
||
cos = cos.to(q.dtype).unsqueeze(0).unsqueeze(0) # [1, 1, HW, head_dim]
|
||
sin = sin.to(q.dtype).unsqueeze(0).unsqueeze(0)
|
||
q_rot = (q * cos) + (self._rotate_half(q) * sin)
|
||
k_rot = (k * cos) + (self._rotate_half(k) * sin)
|
||
return q_rot, k_rot
|
||
|
||
|
||
# ============================================================
|
||
# SE: Squeeze-Excite
|
||
# ============================================================
|
||
|
||
class SqueezeExcite(nn.Module):
|
||
"""Standard Squeeze-Excite channel attention."""
|
||
|
||
def __init__(self, channels: int, reduction: int = 16) -> None:
|
||
super().__init__()
|
||
hidden = max(1, channels // reduction)
|
||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||
self.fc = nn.Sequential(
|
||
nn.Conv2d(channels, hidden, 1),
|
||
nn.SiLU(inplace=True),
|
||
nn.Conv2d(hidden, channels, 1),
|
||
nn.Sigmoid(),
|
||
)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
s = self.pool(x)
|
||
s = self.fc(s)
|
||
return x * s
|
||
|
||
|
||
# ============================================================
|
||
# LayerNorm2d: LN over channels for (B, C, H, W) layout
|
||
# ============================================================
|
||
|
||
class LayerNorm2d(nn.Module):
|
||
"""LayerNorm over C dimension for 4D tensors."""
|
||
|
||
def __init__(self, channels: int, eps: float = 1e-6) -> None:
|
||
super().__init__()
|
||
self.norm = nn.LayerNorm(channels, eps=eps)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
# [B, C, H, W] → [B, H, W, C] → LN → [B, C, H, W]
|
||
x = x.permute(0, 2, 3, 1)
|
||
x = self.norm(x)
|
||
x = x.permute(0, 3, 1, 2).contiguous()
|
||
return x
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# Smoke test
|
||
torch.manual_seed(0)
|
||
|
||
# GGeM test
|
||
g = GGeM(64)
|
||
x = torch.randn(2, 64, 8, 8)
|
||
out = g(x)
|
||
print(f"GGeM out: {out.shape}") # [2, 64]
|
||
|
||
# CHP test — verify rotation invariance
|
||
chp = CircularHarmonicPool(32, rings=8, angles=16, harmonics=4)
|
||
x = torch.randn(1, 32, 16, 16)
|
||
out1 = chp(x)
|
||
# Rotate x by 90° and verify invariance (approximately)
|
||
x_rot = torch.rot90(x, k=1, dims=(-2, -1))
|
||
out2 = chp(x_rot)
|
||
diff = (out1 - out2).abs().max().item()
|
||
print(f"CHP rotation-invariance max diff: {diff:.4e} (should be small)")
|
||
print(f"CHP out shape: {out1.shape}") # [1, 128]
|
||
|
||
# FiLM test
|
||
film = AltitudeFiLM(64)
|
||
x = torch.randn(2, 64, 8, 8)
|
||
altitudes = torch.tensor([150.0, 300.0])
|
||
out = film(x, altitudes)
|
||
print(f"FiLM out: {out.shape}") # [2, 64, 8, 8]
|
||
|
||
# RoPE2D test
|
||
rope = RoPE2D(32)
|
||
q = torch.randn(2, 4, 64, 32) # [B, heads, HW, head_dim]
|
||
k = torch.randn(2, 4, 64, 32)
|
||
q_r, k_r = rope(q, k, 8, 8)
|
||
print(f"RoPE out: q={q_r.shape}, k={k_r.shape}")
|