add sofia models
This commit is contained in:
pikaliov
398
src/models/sofia_v71/layers.py
Normal file
398
src/models/sofia_v71/layers.py
Normal file
@@ -0,0 +1,398 @@
|
||||
"""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}")
|
||||
Reference in New Issue
Block a user