lars
9277d79dff
- Data pipeline: loader (parquet→numpy), transforms (log, local-frame Rodrigues rotation, Normalizer), StepsDataset with event-ID-based split - Model: SinusoidalEmbedding, ConditionEncoder, ResBlock, DenoisingMLP - Schedule: cosine DDPM and conditional flow matching loss (Lipman 2022) - Samplers: flow (Euler ODE), DDPM ancestral, DDIM deterministic - Training loop: AdamW + cosine LR, grad clipping, best-val checkpoint - Validation: per-dimension marginal summary (normalised space) - CLI: TOML config support with CLI-overrides; hyperparam-encoded output directory; config.toml with git hash saved into each run's checkpoint dir - 21 unit tests covering transforms, network, flow/DDPM losses, dataset splits Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class CosineSchedule:
|
|
"""DDPM cosine noise schedule (Nichol & Dhariwal 2021)."""
|
|
|
|
def __init__(self, T: int = 1000, s: float = 0.008) -> None:
|
|
self.T = T
|
|
steps = np.arange(T + 1, dtype=np.float64)
|
|
f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
|
|
alpha_bars = (f / f[0]).astype(np.float32)
|
|
betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)
|
|
|
|
self.betas = torch.from_numpy(betas)
|
|
self.alphas = torch.from_numpy((1.0 - betas))
|
|
self.alpha_bars = torch.from_numpy(alpha_bars[1:])
|
|
|
|
def to(self, device: torch.device) -> "CosineSchedule":
|
|
self.betas = self.betas.to(device)
|
|
self.alphas = self.alphas.to(device)
|
|
self.alpha_bars = self.alpha_bars.to(device)
|
|
return self
|
|
|
|
def q_sample(
|
|
self,
|
|
x0: torch.Tensor,
|
|
t: torch.Tensor,
|
|
noise: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if noise is None:
|
|
noise = torch.randn_like(x0)
|
|
ab = self.alpha_bars[t].view(-1, 1)
|
|
return ab.sqrt() * x0 + (1.0 - ab).sqrt() * noise
|
|
|
|
def loss(
|
|
self,
|
|
model: torch.nn.Module,
|
|
x0: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
B = x0.size(0)
|
|
t = torch.randint(0, self.T, (B,), device=x0.device)
|
|
noise = torch.randn_like(x0)
|
|
x_t = self.q_sample(x0, t, noise)
|
|
t_norm = t.float() / self.T
|
|
pred = model(x_t, t_norm, cond_cont, cond_cat)
|
|
return F.mse_loss(pred, noise)
|
|
|
|
|
|
def flow_matching_loss(
|
|
model: torch.nn.Module,
|
|
x1: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Conditional flow matching loss (Lipman et al. 2022).
|
|
|
|
Straight-line ODE path: x_t = (1-t)*x0 + t*x1, target field u_t = x1-x0.
|
|
"""
|
|
B = x1.size(0)
|
|
t = torch.rand(B, device=x1.device)
|
|
x0 = torch.randn_like(x1)
|
|
x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
|
|
u_t = x1 - x0
|
|
v_t = model(x_t, t, cond_cont, cond_cat)
|
|
return F.mse_loss(v_t, u_t)
|