From 44b0a92e6734cc1954111767bb0261a3e73075ed Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 22 Jul 2026 11:20:28 +0200 Subject: [PATCH] Add WGAN-GP mode as a throwaway fast-eval experiment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds --mode wgan alongside flow/ddpm: both stages get a WGAN-GP generator/critic pair (giant.model.wgan) instead of flow matching, so inference is a single forward pass per stage rather than a 10-step ODE integration — the fast-eval architecture noted in the roadmap. predict/rollout auto-detect the mode from the checkpoint's model_config. Best-checkpoint selection for wgan uses marginal-KL against the EMA generators every epoch, since a critic loss isn't a monotone quality signal. --router is not supported together with --mode wgan. Co-Authored-By: Claude Sonnet 5 --- giant/cli.py | 74 ++++++- giant/config.py | 10 + giant/model/network.py | 260 +++++++++++++++++++++++++ giant/model/wgan.py | 47 +++++ giant/pipeline.py | 25 ++- giant/rollout.py | 26 ++- giant/sample.py | 62 +++++- giant/train.py | 423 +++++++++++++++++++++++++++++++++-------- giant/validate.py | 25 ++- tests/test_wgan.py | 185 ++++++++++++++++++ 10 files changed, 1032 insertions(+), 105 deletions(-) create mode 100644 giant/model/wgan.py create mode 100644 tests/test_wgan.py diff --git a/giant/cli.py b/giant/cli.py index 66d3566..d985efc 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -45,7 +45,12 @@ from giant.model.network import build_models from giant.particles import nearest_known_pdg from giant.pipeline import run_train_job from giant.rollout import rollout as run_rollout -from giant.sample import sample_flow, sample_secondaries +from giant.sample import ( + sample_flow, + sample_secondaries, + sample_wgan, + sample_secondaries_wgan, +) app = typer.Typer(no_args_is_help=True) @@ -177,6 +182,7 @@ def _main() -> None: class Mode(str, Enum): flow = "flow" ddpm = "ddpm" + wgan = "wgan" class Conditioning(str, Enum): @@ -308,6 +314,38 @@ def train( "--router-axis 'pdg:n_experts=3,emb_dim=8'", ), ] = None, + n_critic: Annotated[ + Optional[int], + typer.Option( + "--n-critic", + help="WGAN-GP (--mode wgan only): critic updates per generator " + "update (default: 5)", + ), + ] = None, + gp_weight: Annotated[ + Optional[float], + typer.Option( + "--gp-weight", + help="WGAN-GP (--mode wgan only): gradient-penalty coefficient " + "(default: 10.0)", + ), + ] = None, + noise_dim: Annotated[ + Optional[int], + typer.Option( + "--noise-dim", + help="WGAN (--mode wgan only): generator input noise-vector " + "width (default: 64)", + ), + ] = None, + critic_lr: Annotated[ + Optional[float], + typer.Option( + "--critic-lr", + help="WGAN-GP (--mode wgan only): critic learning rate " + "(default: same as --lr)", + ), + ] = None, val_fraction: Annotated[ Optional[float], typer.Option("--val-fraction", "-f") ] = None, @@ -395,6 +433,9 @@ def train( "validate_every": validate_every, "validate_steps": validate_steps, "max_val_batches": max_val_batches, + "n_critic": n_critic, + "gp_weight": gp_weight, + "critic_lr": critic_lr, }.items() if v is not None } @@ -406,6 +447,7 @@ def train( "emb_dim": emb_dim, "dropout": dropout, "conditioning": conditioning.value if conditioning is not None else None, + "noise_dim": noise_dim, }.items() if v is not None } @@ -510,7 +552,12 @@ def predict( ), ] = "4096", steps: Annotated[ - int, typer.Option("--steps", "-s", help="Flow matching ODE steps") + int, + typer.Option( + "--steps", + "-s", + help="Flow matching ODE steps (ignored for a wgan checkpoint)", + ), ] = 10, weights: Annotated[ Weights, @@ -649,12 +696,20 @@ def predict( cc = torch.from_numpy(cond_cont).float().to(_device) ck = torch.from_numpy(cond_cat).long().to(_device) - stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps) + if model_cfg.get("mode") == "wgan": + stage1_norm, n_sec_pred = sample_wgan(model, cc, ck) + else: + stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps) if coord == Coord.global_: - sec_cont, sec_phys, _sec_valid_pred = sample_secondaries( - sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps - ) + if model_cfg.get("mode") == "wgan": + sec_cont, sec_phys, _sec_valid_pred = sample_secondaries_wgan( + sec_decoder, cc, ck, stage1_norm, n_sec_pred + ) + else: + sec_cont, sec_phys, _sec_valid_pred = sample_secondaries( + sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps + ) sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy() n_sec_pred_np = n_sec_pred.cpu().numpy() @@ -898,7 +953,11 @@ def rollout( ] = 1000, steps: Annotated[ int, - typer.Option("--steps", "-s", help="Flow matching ODE steps per model call"), + typer.Option( + "--steps", + "-s", + help="Flow matching ODE steps per model call (ignored for a wgan checkpoint)", + ), ] = 10, weights: Annotated[ Weights, @@ -1030,6 +1089,7 @@ def rollout( escape_threshold=escape_threshold, on_chunk=_write_chunk, conditioning=conditioning, + mode=model_cfg.get("mode", "flow"), ) if writer is not None: writer.close() diff --git a/giant/config.py b/giant/config.py index ab121eb..cfd7a1b 100644 --- a/giant/config.py +++ b/giant/config.py @@ -27,12 +27,22 @@ DEFAULT_CONFIG: dict = { "warmup_epochs": 5, "lambda_nsec": 0.1, "lambda_s2": 1.0, + # WGAN-GP-only knobs (mode == "wgan"; ignored by flow/ddpm). n_critic: + # critic updates per generator update. gp_weight: gradient-penalty + # coefficient (Gulrajani et al. 2017). critic_lr: 0.0 means "use + # `lr`" — not None, since save_config's TOML writer has no null + # literal to round-trip. + "n_critic": 5, + "gp_weight": 10.0, + "critic_lr": 0.0, }, "model": { "hidden_dim": 256, "n_blocks": 6, "emb_dim": 16, "dropout": 0.1, + # WGAN generator noise-vector width (mode == "wgan" only). + "noise_dim": 64, # "physical" conditions on material/particle physical properties via # a small MLP (giant.model.network.ConditionEncoder); "embedding" # keeps the original learned pdg/material embedding tables — kept diff --git a/giant/model/network.py b/giant/model/network.py index 9646a0c..5596a17 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -297,6 +297,233 @@ class SecondaryDecoder(nn.Module): return self.out_proj(x) +class WGANGenerator(nn.Module): + """Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time. + + Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the + input is a noise vector `z` (not a diffused/interpolated `x_t`) and the + ResBlocks condition on the condition encoding alone (no time embedding to + concatenate) — see `giant/model/wgan.py` for the adversarial losses, and + `giant.sample.sample_wgan` for single-pass sampling. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + noise_dim: int = 64, + dropout: float = 0.1, + k_max: int = K_MAX, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.noise_dim = noise_dim + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(noise_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, cond_out_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_proj = nn.Linear(hidden_dim, x_dim) + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max + 1), + ) + + def forward( + self, + z: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat) + x = self.input_proj(z) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + def predict_n_sec( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone.""" + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + +class Critic(nn.Module): + """Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`. + + Kept structurally parallel to `WGANGenerator` (own condition encoder — + separate weights from the generator's, standard GAN practice) but has no + n_sec head: n_sec is never adversarial, it stays a plain classifier on + the generator side. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(x_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, cond_out_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_norm = nn.LayerNorm(hidden_dim) + self.out_proj = nn.Linear(hidden_dim, 1) + + def forward( + self, + x: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat) + h = self.input_proj(x) + for block in self.blocks: + h = block(h, cond) + return self.out_proj(self.out_norm(h)).squeeze(-1) + + +class WGANSecondaryGenerator(nn.Module): + """Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots. + + Mirrors `SecondaryDecoder` minus the time embedding, the same way + `WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`, + conditions on `SecondaryConditionEncoder`'s output alone. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + noise_dim: int = 64, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.noise_dim = noise_dim + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(noise_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, cond_out_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_proj = nn.Linear(hidden_dim, sec_dim) + + def forward( + self, + z: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat, stage1_out) + x = self.input_proj(z) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +class SecondaryCritic(nn.Module): + """Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots.""" + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(sec_dim, hidden_dim) + self.blocks = nn.ModuleList( + [ + ResBlock(hidden_dim, cond_out_dim, dropout=dropout) + for _ in range(n_blocks) + ] + ) + self.out_norm = nn.LayerNorm(hidden_dim) + self.out_proj = nn.Linear(hidden_dim, 1) + + def forward( + self, + x: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat, stage1_out) + h = self.input_proj(x) + for block in self.blocks: + h = block(h, cond) + return self.out_proj(self.out_norm(h)).squeeze(-1) + + class Router(nn.Module): """Contract for a pluggable mixture-of-experts routing axis. @@ -799,6 +1026,11 @@ _SEC_DECODER_MODEL_KEYS = { "dropout", "conditioning", } +_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"} +_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"} +# Critic has no n_sec head (n_sec is never adversarial), so it doesn't accept +# k_max the way DenoisingMLP/WGANGenerator do. +_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"} _AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") @@ -864,6 +1096,19 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: once and passed to both stage1/sec_decoder, so they structurally always share one mode. """ + if model_config.get("mode") == "wgan": + stage1 = WGANGenerator( + **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS} + ) + sec_decoder = WGANSecondaryGenerator( + **{ + k: v + for k, v in model_config.items() + if k in _WGAN_SEC_GENERATOR_MODEL_KEYS + } + ) + return stage1, sec_decoder + router_cfg = model_config.get("router") if router_cfg and router_cfg.get("enabled"): pdg_vocab = model_config["pdg_vocab"] @@ -895,3 +1140,18 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS} ) return stage1, sec_decoder + + +def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]: + """Construct (critic, sec_critic) for `--mode wgan` training. + + Training-only — never persisted for inference the way `build_models`'s + pair is, since `predict`/`rollout` only ever run the generators. + """ + critic = Critic( + **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS} + ) + sec_critic = SecondaryCritic( + **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS} + ) + return critic, sec_critic diff --git a/giant/model/wgan.py b/giant/model/wgan.py new file mode 100644 index 0000000..161b416 --- /dev/null +++ b/giant/model/wgan.py @@ -0,0 +1,47 @@ +from typing import Callable + +import torch + +CriticFn = Callable[[torch.Tensor], torch.Tensor] + + +def gradient_penalty( + critic_fn: CriticFn, + real: torch.Tensor, + fake: torch.Tensor, + mask: torch.Tensor | None = None, +) -> torch.Tensor: + """WGAN-GP penalty (Gulrajani et al. 2017): (||grad||_2 - 1)^2 at a random interpolate. + + `mask` (same shape as `real`/`fake`, 1 for real content / 0 for padding) + is for Stage 2's variable-length slot vector: both the interpolate and the + critic's gradient are zeroed on padded dims first, so the norm target of 1 + is only ever asked of genuine content, not the padding convention shared + by both `real` and `fake`. + """ + eps = torch.rand(real.size(0), 1, device=real.device) + x_hat = eps * real + (1 - eps) * fake + if mask is not None: + x_hat = x_hat * mask + x_hat = x_hat.requires_grad_(True) + scores = critic_fn(x_hat) + grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0] + if mask is not None: + grad = grad * mask + return ((grad.norm(2, dim=1) - 1) ** 2).mean() + + +def critic_loss( + critic_fn: CriticFn, + real: torch.Tensor, + fake: torch.Tensor, + gp_weight: float, + mask: torch.Tensor | None = None, +) -> torch.Tensor: + """WGAN-GP critic loss. `fake` must already be `.detach()`'d by the caller.""" + gp = gradient_penalty(critic_fn, real, fake, mask=mask) + return critic_fn(fake).mean() - critic_fn(real).mean() + gp_weight * gp + + +def generator_loss(critic_fn: CriticFn, fake: torch.Tensor) -> torch.Tensor: + return -critic_fn(fake).mean() diff --git a/giant/pipeline.py b/giant/pipeline.py index f3db978..c06a0d0 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -22,7 +22,7 @@ from giant.data.loader import ( ) from giant.data.transforms import build_features, _WelfordAccumulator from giant.data.dataset import make_event_split, StreamingStepsDataset -from giant.model.network import build_models +from giant.model.network import build_models, build_critics from giant.train import train as run_training @@ -63,6 +63,11 @@ def run_train_job( echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials") router_cfg = m["router"] + if t["mode"] == "wgan" and router_cfg.get("enabled"): + raise ValueError( + "--mode wgan does not support --router (no routed WGAN generator/" + "critic exists) — disable one or the other" + ) proc_map: dict[str, int] | None = None if router_cfg.get("enabled") and router_cfg.get("type") == "process": echo("building process vocabulary …") @@ -159,6 +164,10 @@ def run_train_job( "router": dict(router_cfg), "expert_hidden_dim": router_cfg["expert_hidden_dim"], "expert_n_blocks": router_cfg["expert_n_blocks"], + # Read by `predict`/`rollout` (which never receive their own --mode + # flag) to auto-detect which sampler a checkpoint needs. + "mode": t["mode"], + "noise_dim": m.get("noise_dim", 64), } stage1_model, sec_decoder = build_models(model_config) @@ -167,6 +176,15 @@ def run_train_job( f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters" ) + critic = None + sec_critic = None + if t["mode"] == "wgan": + critic, sec_critic = build_critics(model_config) + echo( + f"critic: {sum(p.numel() for p in critic.parameters()):,} parameters | " + f"sec_critic: {sum(p.numel() for p in sec_critic.parameters()):,} parameters" + ) + out_dir.mkdir(parents=True, exist_ok=True) meta = config.build_run_meta( data=data, @@ -210,4 +228,9 @@ def run_train_job( validate_steps=t["validate_steps"], max_val_batches=t["max_val_batches"], total_train_batches=total_train_batches, + critic=critic, + sec_critic=sec_critic, + n_critic=t.get("n_critic", 5), + gp_weight=t.get("gp_weight", 10.0), + critic_lr=t.get("critic_lr") or None, ) diff --git a/giant/rollout.py b/giant/rollout.py index 0acf148..5fde841 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -39,7 +39,12 @@ from giant.data.transforms import ( reconstruct_post_pos, ) from giant.particles import nearest_known_pdg, particle_phys_array -from giant.sample import sample_flow, sample_secondaries +from giant.sample import ( + sample_flow, + sample_secondaries, + sample_wgan, + sample_secondaries_wgan, +) # Record columns produced per step / per terminal marker. _RECORD_KEYS = [ @@ -306,6 +311,7 @@ def rollout( escape_threshold: float | None = None, on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, conditioning: str = "embedding", + mode: str = "flow", ) -> dict[str, np.ndarray] | RolloutSummary: """Run showers to completion. @@ -359,6 +365,7 @@ def rollout( device, max_tracks_per_event, conditioning, + mode, ) ) frontier = _concat_frontiers(next_parts) @@ -389,6 +396,7 @@ def _step_chunk( device, max_tracks_per_event, conditioning, + mode="flow", ) -> dict[str, np.ndarray]: """Advance one chunk of tracks by a single step; return the next frontier.""" n = len(tr["event_id"]) @@ -463,7 +471,10 @@ def _step_chunk( cc = torch.from_numpy(cond_cont).float().to(device) ck = torch.from_numpy(cond_cat).long().to(device) - stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps) + if mode == "wgan": + stage1_norm, n_sec_pred = sample_wgan(stage1_model, cc, ck) + else: + stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps) raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy()) step_length = inv_log_transform(raw[:, 0]) @@ -490,9 +501,14 @@ def _step_chunk( # identity, used as-is for the spawned track's own future conditioning. # sec_pdg_code below is a *separate*, reporting-only nearest-known-PDG # label (never fed back into the model) — see giant/particles.py. - sec_cont, sec_phys, _valid = sample_secondaries( - sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps - ) + if mode == "wgan": + sec_cont, sec_phys, _valid = sample_secondaries_wgan( + sec_decoder, cc, ck, stage1_norm, n_sec_pred + ) + else: + sec_cont, sec_phys, _valid = sample_secondaries( + sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps + ) sec_full = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy() sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries( sec_full, diff --git a/giant/sample.py b/giant/sample.py index e97fead..271897b 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -3,6 +3,26 @@ import torch from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM +def _slots_from_flat( + x: torch.Tensor, n_sec_pred: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reshape a flat (B, SEC_DIM) decoder output into per-slot tensors. + + Returns (sec_cont, sec_phys, sec_valid) — see `sample_secondaries`'s + docstring for their shapes/meaning. Shared by both the flow-matching and + WGAN Stage-2 samplers, which differ only in how `x` was produced. + """ + B = x.size(0) + device = x.device + x_slots = x.view(B, K_MAX, SEC_SLOT_DIM) + sec_cont = x_slots[:, :, :4] + sec_phys = x_slots[:, :, 4:] + sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze( + 1 + ) + return sec_cont, sec_phys, sec_valid + + @torch.no_grad() def sample_flow( model: torch.nn.Module, @@ -64,14 +84,7 @@ def sample_secondaries( v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out) x = x + v * dt - x_slots = x.view(B, K_MAX, SEC_SLOT_DIM) - sec_cont = x_slots[:, :, :4] - sec_phys = x_slots[:, :, 4:] - - sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze( - 1 - ) - return sec_cont, sec_phys, sec_valid + return _slots_from_flat(x, n_sec_pred) @torch.no_grad() @@ -102,6 +115,39 @@ def sample_ddpm( return x, n_sec_pred +@torch.no_grad() +def sample_wgan( + generator: torch.nn.Module, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred).""" + generator.eval() + B = cond_cont.size(0) + z = torch.randn(B, generator.noise_dim, device=cond_cont.device) + x = generator(z, cond_cont, cond_cat) + n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat) + n_sec_pred = n_sec_logits.argmax(dim=-1) + return x, n_sec_pred + + +@torch.no_grad() +def sample_secondaries_wgan( + sec_decoder: torch.nn.Module, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + n_sec_pred: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Single-pass Stage-2 WGAN generator sample; see `sample_secondaries`'s + docstring for the returned (sec_cont, sec_phys, sec_valid) shapes.""" + sec_decoder.eval() + B = cond_cont.size(0) + z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device) + x = sec_decoder(z, cond_cont, cond_cat, stage1_out) + return _slots_from_flat(x, n_sec_pred) + + @torch.no_grad() def sample_ddim( model: torch.nn.Module, diff --git a/giant/train.py b/giant/train.py index 1849eaa..02dda14 100644 --- a/giant/train.py +++ b/giant/train.py @@ -8,17 +8,20 @@ from pathlib import Path from types import FrameType from typing import Callable +import numpy as np import torch import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm +from giant.constants import K_MAX, SEC_SLOT_DIM from giant.model.schedule import ( CosineSchedule, flow_matching_loss, flow_matching_loss_secondary, ) +from giant.model.wgan import gradient_penalty, generator_loss from giant.validate import validate_marginals _METRICS_FIELDS = [ @@ -29,12 +32,17 @@ _METRICS_FIELDS = [ "train_loss_s2", "train_loss_balance", "train_loss_proc", + "d_loss", + "g_loss", + "wasserstein_estimate", + "gp_loss", "val_loss", "val_loss_s1", "val_loss_nsec", "val_loss_s2", "val_loss_balance", "val_loss_proc", + "val_marginal_kl", "lr", "grad_norm", "epoch_time_s", @@ -129,8 +137,6 @@ def _compute_losses( # giant.data.transforms.encode_secondaries) rather than a learned/moving # one, so — unlike the embedding-table target this replaced — no # detaching is needed to keep the target from chasing the decoder. - from giant.constants import K_MAX - x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) @@ -168,6 +174,120 @@ def _compute_losses( return total, l_s1, l_nsec, l_s2, l_balance, l_proc +def _wgan_train_step( + generator: torch.nn.Module, + sec_generator: torch.nn.Module, + critic: torch.nn.Module, + sec_critic: torch.nn.Module, + batch: tuple, + device: torch.device, + optimizer_g: optim.Optimizer, + optimizer_d: optim.Optimizer, + g_params: list, + d_params: list, + step_count: int, + n_critic: int, + gp_weight: float, + lambda_nsec: float, + lambda_s2: float, +) -> dict: + """One WGAN-GP training step, both stages (see giant/model/wgan.py for the losses). + + Both critics update every batch. Every `n_critic`-th batch additionally + updates both generators. The (non-adversarial) n_sec classifier updates + every batch regardless — folded into whichever `optimizer_g` step happens + this batch (full adversarial g_loss on generator batches, n_sec-only in + between) rather than throttled to the generator's cadence, since n_sec + accuracy is a headline flow-vs-wgan comparison metric and shares the + generator's ConditionEncoder. + + Stage 2's real/fake target is a flattened (B, SEC_DIM) vector with + `K_MAX - n_sec` padded slots per row; both critic's input and its + gradient-penalty gradient are masked to the valid slots (see + `giant.model.wgan.gradient_penalty`) so the critic can't key on padding + instead of genuine content. Stage 2 is conditioned on the *real* + ground-truth Stage-1 target (`x1_s1`, detached) rather than the + generator's own fake Stage-1 output — same precedent as the flow-matching + path's `flow_matching_loss_secondary` call, avoiding compounding errors + during training. + """ + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = batch + cond_cont = cond_cont.to(device) + cond_cat = cond_cat.to(device) + x1_s1 = x1_s1.to(device) + n_sec = n_sec.to(device) + sec_cont = sec_cont.to(device) + + B = x1_s1.size(0) + x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) + sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) + mask_flat = ( + sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float() + ) + stage1_ctx = x1_s1.detach() + + def critic_fn1(x: torch.Tensor) -> torch.Tensor: + return critic(x, cond_cont, cond_cat) + + def critic_fn2(x: torch.Tensor) -> torch.Tensor: + return sec_critic(x, cond_cont, cond_cat, stage1_ctx) + + z1 = torch.randn(B, generator.noise_dim, device=device) + fake1 = generator(z1, cond_cont, cond_cat) + z2 = torch.randn(B, sec_generator.noise_dim, device=device) + fake2 = sec_generator(z2, cond_cont, cond_cat, stage1_ctx) + fake2_masked = fake2 * mask_flat + real2_masked = x1_s2 * mask_flat + + # --- Critic step (every batch) --- + fake1_detached = fake1.detach() + real1_score = critic_fn1(x1_s1) + fake1_score = critic_fn1(fake1_detached) + gp1 = gradient_penalty(critic_fn1, x1_s1, fake1_detached) + d1 = fake1_score.mean() - real1_score.mean() + gp_weight * gp1 + wasserstein_estimate = (real1_score.mean() - fake1_score.mean()).detach() + + fake2_detached_masked = fake2_masked.detach() + real2_score = critic_fn2(real2_masked) + fake2_score = critic_fn2(fake2_detached_masked) + gp2 = gradient_penalty( + critic_fn2, real2_masked, fake2_detached_masked, mask=mask_flat + ) + d2 = fake2_score.mean() - real2_score.mean() + gp_weight * gp2 + + d_loss = d1 + lambda_s2 * d2 + optimizer_d.zero_grad() + d_loss.backward() + grad_norm_d = torch.nn.utils.clip_grad_norm_(d_params, 1.0) + optimizer_d.step() + + # --- Generator (+ n_sec) step --- + did_g_step = step_count % n_critic == 0 + l_nsec = F.cross_entropy(generator.predict_n_sec(cond_cont, cond_cat), n_sec) + optimizer_g.zero_grad() + if did_g_step: + g1 = generator_loss(critic_fn1, fake1) + g2 = generator_loss(critic_fn2, fake2_masked) + g_loss = g1 + lambda_nsec * l_nsec + lambda_s2 * g2 + else: + g1 = torch.zeros((), device=device) + g2 = torch.zeros((), device=device) + g_loss = lambda_nsec * l_nsec + g_loss.backward() + grad_norm_g = torch.nn.utils.clip_grad_norm_(g_params, 1.0) + optimizer_g.step() + + return { + "d_loss": d_loss.detach(), + "g_loss": (g1 + lambda_s2 * g2).detach(), + "wasserstein_estimate": wasserstein_estimate, + "gp_loss": (gp1 + lambda_s2 * gp2).detach(), + "l_nsec": l_nsec.detach(), + "did_g_step": did_g_step, + "grad_norm": grad_norm_d.item() + grad_norm_g.item(), + } + + def train( stage1_model: torch.nn.Module, sec_decoder: torch.nn.Module, @@ -195,12 +315,23 @@ def train( validate_steps: int = 10, max_val_batches: int = 0, total_train_batches: int = 0, + critic: torch.nn.Module | None = None, + sec_critic: torch.nn.Module | None = None, + n_critic: int = 5, + gp_weight: float = 10.0, + critic_lr: float | None = None, ) -> None: out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) stage1_model = stage1_model.to(device) sec_decoder = sec_decoder.to(device) + if mode == "wgan": + assert critic is not None and sec_critic is not None, ( + "mode='wgan' requires critic/sec_critic (see giant.model.network.build_critics)" + ) + critic = critic.to(device) + sec_critic = sec_critic.to(device) # Flow-matching/diffusion models sample noticeably better from an EMA of # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed @@ -217,14 +348,36 @@ def train( p.requires_grad_(False) all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) - optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay) + all_params_d: list = [] + optimizer_d: optim.Optimizer | None = None + if mode == "wgan": + assert critic is not None and sec_critic is not None + # Standard WGAN-GP recipe (Gulrajani et al. 2017): Adam with + # beta1=0 (momentum destabilizes critic training) and no weight + # decay, rather than the AdamW(weight_decay=...) used for flow/ddpm. + optimizer = optim.Adam(all_params, lr=lr, betas=(0.0, 0.9)) + all_params_d = list(critic.parameters()) + list(sec_critic.parameters()) + optimizer_d = optim.Adam( + all_params_d, + lr=critic_lr if critic_lr is not None else lr, + betas=(0.0, 0.9), + ) + else: + optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay) # Warmup/decay in units of optimizer steps rather than epochs: at large # dataset sizes a single epoch can be tens of thousands of steps, and an # epoch-granularity schedule would leave warmup/cosine decay unable to # move within it. Requires an accurate `total_train_batches` (steps per # epoch); the only caller, run_train_job, always supplies one. + # + # In wgan mode, `lr_sched.step()`/EMA only fire on generator steps (see + # the per-batch loop below) — 1 in every `n_critic` batches — so the + # schedule's own step-counting must be in those same units, or warmup + # would never finish and cosine decay would barely move. steps_per_epoch = max(total_train_batches, 1) + if mode == "wgan": + steps_per_epoch = max(total_train_batches // (n_critic + 1), 1) warmup_steps = warmup_epochs * steps_per_epoch total_steps = max(epochs * steps_per_epoch, 1) @@ -251,6 +404,15 @@ def train( ema_sec_decoder.load_state_dict( ckpt.get("sec_decoder_ema", ckpt["sec_decoder"]) ) + if mode == "wgan": + assert ( + critic is not None + and sec_critic is not None + and optimizer_d is not None + ) + critic.load_state_dict(ckpt["critic"]) + sec_critic.load_state_dict(ckpt["sec_critic"]) + optimizer_d.load_state_dict(ckpt["optimizer_d"]) optimizer.load_state_dict(ckpt["optimizer"]) lr_sched.load_state_dict(ckpt["lr_sched"]) start_epoch = ckpt.get("epoch", 0) + 1 @@ -284,17 +446,26 @@ def train( epoch_w = len(str(epochs)) last_completed_epoch = start_epoch - 1 + global_step = 0 with _GracefulShutdown() as shutdown: for epoch in range(start_epoch, epochs + 1): epoch_start = time.monotonic() stage1_model.train() sec_decoder.train() + if mode == "wgan": + assert critic is not None and sec_critic is not None + critic.train() + sec_critic.train() train_loss_sum = 0.0 train_s1_sum = 0.0 train_nsec_sum = 0.0 train_s2_sum = 0.0 train_balance_sum = 0.0 train_proc_sum = 0.0 + train_d_sum = 0.0 + train_g_sum = 0.0 + train_wasserstein_sum = 0.0 + train_gp_sum = 0.0 train_n = 0 train_batches = 0 grad_norm_sum = 0.0 @@ -309,37 +480,84 @@ def train( dynamic_ncols=True, ) for batch in bar: - loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( - stage1_model, - sec_decoder, - batch, - mode, - ddpm_schedule, - device, - lambda_nsec, - lambda_s2, - lambda_balance, - lambda_proc, - ) - optimizer.zero_grad() - loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0) - optimizer.step() - lr_sched.step() - if ema_decay > 0: - assert ema_stage1_model is not None and ema_sec_decoder is not None - _update_ema(ema_stage1_model, stage1_model, ema_decay) - _update_ema(ema_sec_decoder, sec_decoder, ema_decay) + if mode == "wgan": + assert ( + critic is not None + and sec_critic is not None + and optimizer_d is not None + ) + stats = _wgan_train_step( + stage1_model, + sec_decoder, + critic, + sec_critic, + batch, + device, + optimizer, + optimizer_d, + all_params, + all_params_d, + global_step, + n_critic, + gp_weight, + lambda_nsec, + lambda_s2, + ) + global_step += 1 + if stats["did_g_step"]: + lr_sched.step() + if ema_decay > 0: + assert ( + ema_stage1_model is not None + and ema_sec_decoder is not None + ) + _update_ema(ema_stage1_model, stage1_model, ema_decay) + _update_ema(ema_sec_decoder, sec_decoder, ema_decay) + + B = batch[0].size(0) + batch_loss = stats["d_loss"].item() + stats["g_loss"].item() + batch_grad_norm = stats["grad_norm"] + train_loss_sum += batch_loss * B + train_nsec_sum += stats["l_nsec"].item() * B + train_d_sum += stats["d_loss"].item() * B + train_g_sum += stats["g_loss"].item() * B + train_wasserstein_sum += stats["wasserstein_estimate"].item() * B + train_gp_sum += stats["gp_loss"].item() * B + else: + loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( + stage1_model, + sec_decoder, + batch, + mode, + ddpm_schedule, + device, + lambda_nsec, + lambda_s2, + lambda_balance, + lambda_proc, + ) + optimizer.zero_grad() + loss.backward() + grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0) + optimizer.step() + lr_sched.step() + if ema_decay > 0: + assert ( + ema_stage1_model is not None and ema_sec_decoder is not None + ) + _update_ema(ema_stage1_model, stage1_model, ema_decay) + _update_ema(ema_sec_decoder, sec_decoder, ema_decay) + + B = batch[0].size(0) + batch_loss = loss.item() + batch_grad_norm = grad_norm.item() + train_loss_sum += batch_loss * B + train_s1_sum += l_s1.item() * B + train_nsec_sum += l_nsec.item() * B + train_s2_sum += l_s2.item() * B + train_balance_sum += l_balance.item() * B + train_proc_sum += l_proc.item() * B - B = batch[0].size(0) - batch_loss = loss.item() - batch_grad_norm = grad_norm.item() - train_loss_sum += batch_loss * B - train_s1_sum += l_s1.item() * B - train_nsec_sum += l_nsec.item() * B - train_s2_sum += l_s2.item() * B - train_balance_sum += l_balance.item() * B - train_proc_sum += l_proc.item() * B train_n += B train_batches += 1 grad_norm_sum += batch_grad_norm @@ -368,38 +586,87 @@ def train( stage1_model.eval() sec_decoder.eval() - val_loss_sum = 0.0 - val_s1_sum = 0.0 - val_nsec_sum = 0.0 - val_s2_sum = 0.0 - val_balance_sum = 0.0 - val_proc_sum = 0.0 - val_n = 0 - with torch.no_grad(): - for val_batch_idx, batch in enumerate(val_loader): - if max_val_batches > 0 and val_batch_idx >= max_val_batches: - break - loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( + if mode == "wgan": + assert critic is not None and sec_critic is not None + critic.eval() + sec_critic.eval() + + val_marginal_kl = float("nan") + if mode == "wgan": + # WGANGenerator.forward(z, cond_cont, cond_cat) has no + # diffusion/flow `t` argument, so the usual _compute_losses + # val loop below (which calls flow_matching_loss -> + # stage1_model(x_t, t, ...)) doesn't apply — and a WGAN + # critic loss isn't a monotone quality signal fit for + # best-checkpoint selection anyway. Select on marginal KL + # against the EMA generators instead (matches what + # predict/rollout sample from by default, --weights ema). + eval_stage1 = ( + ema_stage1_model if ema_stage1_model is not None else stage1_model + ) + eval_sec_decoder = ( + ema_sec_decoder if ema_sec_decoder is not None else sec_decoder + ) + marginal_result = validate_marginals( + eval_stage1, + val_loader, + mode=mode, + device=device, + sec_decoder=eval_sec_decoder, + ) + val_marginal_kl = float(np.mean(marginal_result["kl_divergence"])) + val_loss = val_marginal_kl + val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = ( + val_proc_sum + ) = 0.0 + val_n = 1 + else: + val_loss_sum = 0.0 + val_s1_sum = 0.0 + val_nsec_sum = 0.0 + val_s2_sum = 0.0 + val_balance_sum = 0.0 + val_proc_sum = 0.0 + val_n = 0 + with torch.no_grad(): + for val_batch_idx, batch in enumerate(val_loader): + if max_val_batches > 0 and val_batch_idx >= max_val_batches: + break + loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses( + stage1_model, + sec_decoder, + batch, + mode, + ddpm_schedule, + device, + lambda_nsec, + lambda_s2, + lambda_balance, + lambda_proc, + ) + B = batch[0].size(0) + val_loss_sum += loss.item() * B + val_s1_sum += l_s1.item() * B + val_nsec_sum += l_nsec.item() * B + val_s2_sum += l_s2.item() * B + val_balance_sum += l_balance.item() * B + val_proc_sum += l_proc.item() * B + val_n += B + val_loss = val_loss_sum / max(val_n, 1) + + if validate_every > 0 and epoch % validate_every == 0: + print(f"[epoch {epoch}] marginal validation:") + marginal_result = validate_marginals( stage1_model, - sec_decoder, - batch, - mode, - ddpm_schedule, - device, - lambda_nsec, - lambda_s2, - lambda_balance, - lambda_proc, + val_loader, + mode=mode, + schedule=ddpm_schedule, + device=device, + steps=validate_steps, + sec_decoder=sec_decoder, ) - B = batch[0].size(0) - val_loss_sum += loss.item() * B - val_s1_sum += l_s1.item() * B - val_nsec_sum += l_nsec.item() * B - val_s2_sum += l_s2.item() * B - val_balance_sum += l_balance.item() * B - val_proc_sum += l_proc.item() * B - val_n += B - val_loss = val_loss_sum / max(val_n, 1) + val_marginal_kl = float(np.mean(marginal_result["kl_divergence"])) + epoch_time = time.monotonic() - epoch_start is_best = val_loss < best_val_loss @@ -411,7 +678,9 @@ def train( f" nsec={train_nsec_sum / max(train_n, 1):.3f}" f" s2={train_s2_sum / max(train_n, 1):.3f}" f" bal={train_balance_sum / max(train_n, 1):.3f}" - f" proc={train_proc_sum / max(train_n, 1):.3f})" + f" proc={train_proc_sum / max(train_n, 1):.3f}" + f" d={train_d_sum / max(train_n, 1):.3f}" + f" g={train_g_sum / max(train_n, 1):.3f})" f" val {val_loss:.4f}" f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}" f" {epoch_time:.1f}s{marker}" @@ -425,12 +694,17 @@ def train( "train_loss_s2": train_s2_sum / max(train_n, 1), "train_loss_balance": train_balance_sum / max(train_n, 1), "train_loss_proc": train_proc_sum / max(train_n, 1), + "d_loss": train_d_sum / max(train_n, 1), + "g_loss": train_g_sum / max(train_n, 1), + "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1), + "gp_loss": train_gp_sum / max(train_n, 1), "val_loss": val_loss, "val_loss_s1": val_s1_sum / max(val_n, 1), "val_loss_nsec": val_nsec_sum / max(val_n, 1), "val_loss_s2": val_s2_sum / max(val_n, 1), "val_loss_balance": val_balance_sum / max(val_n, 1), "val_loss_proc": val_proc_sum / max(val_n, 1), + "val_marginal_kl": val_marginal_kl, "lr": current_lr, "grad_norm": train_grad_norm, "epoch_time_s": epoch_time, @@ -438,18 +712,6 @@ def train( ) metrics_file.flush() - if validate_every > 0 and epoch % validate_every == 0: - print(f"[epoch {epoch}] marginal validation:") - validate_marginals( - stage1_model, - val_loader, - mode=mode, - schedule=ddpm_schedule, - device=device, - steps=validate_steps, - sec_decoder=sec_decoder, - ) - ckpt: dict = { "model": stage1_model.state_dict(), "sec_decoder": sec_decoder.state_dict(), @@ -458,6 +720,15 @@ def train( "epoch": epoch, "best_val_loss": best_val_loss, } + if mode == "wgan": + assert ( + critic is not None + and sec_critic is not None + and optimizer_d is not None + ) + ckpt["critic"] = critic.state_dict() + ckpt["sec_critic"] = sec_critic.state_dict() + ckpt["optimizer_d"] = optimizer_d.state_dict() if ema_decay > 0: assert ema_stage1_model is not None and ema_sec_decoder is not None ckpt["model_ema"] = ema_stage1_model.state_dict() diff --git a/giant/validate.py b/giant/validate.py index ae69207..44ae071 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -8,6 +8,8 @@ from giant.sample import ( sample_ddpm, sample_ddim, sample_secondaries, + sample_wgan, + sample_secondaries_wgan, ) _SEC_PHYS_NAMES = ["log_mass", "charge"] @@ -98,6 +100,8 @@ def validate_marginals( gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps)) elif mode == "ddpm": gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule) + elif mode == "wgan": + gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat) else: gen, n_sec_pred = sample_ddim( model, cond_cont, cond_cat, schedule, **_kw(steps) @@ -118,14 +122,19 @@ def validate_marginals( real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64))) real_phys = sec_cont[:, :, 4:6].numpy() # (B, K_MAX, 2) [log_mass, charge] - sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries( - sec_decoder, - cond_cont, - cond_cat, - gen, - n_sec_pred, - steps=steps if steps is not None else 10, - ) + if mode == "wgan": + sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan( + sec_decoder, cond_cont, cond_cat, gen, n_sec_pred + ) + else: + sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries( + sec_decoder, + cond_cont, + cond_cat, + gen, + n_sec_pred, + steps=steps if steps is not None else 10, + ) gen_frac = 1.0 / ( 1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)) ) diff --git a/tests/test_wgan.py b/tests/test_wgan.py new file mode 100644 index 0000000..4b85b0c --- /dev/null +++ b/tests/test_wgan.py @@ -0,0 +1,185 @@ +import torch + +from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM +from giant.model.network import ( + Critic, + SecondaryCritic, + WGANGenerator, + WGANSecondaryGenerator, +) +from giant.model.wgan import critic_loss, generator_loss, gradient_penalty +from giant.sample import sample_secondaries_wgan, sample_wgan + + +def _cond(B=8): + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + return cond_cont, cond_cat + + +def _small_generator(): + return WGANGenerator( + pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8 + ) + + +def _small_critic(): + return Critic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2) + + +def _small_sec_generator(): + return WGANSecondaryGenerator( + pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8 + ) + + +def _small_sec_critic(): + return SecondaryCritic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2) + + +def _mask(B, n_sec): + sec_mask = torch.arange(K_MAX).unsqueeze(0) < n_sec.unsqueeze(1) + return sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float() + + +# --- Stage-1 generator/critic --- + + +def test_wgan_generator_output_shape(): + B = 8 + model = _small_generator() + cond_cont, cond_cat = _cond(B) + z = torch.randn(B, model.noise_dim) + out = model(z, cond_cont, cond_cat) + assert out.shape == (B, X_DIM) + + +def test_wgan_generator_predict_n_sec_shape(): + B = 6 + model = _small_generator() + cond_cont, cond_cat = _cond(B) + logits = model.predict_n_sec(cond_cont, cond_cat) + assert logits.shape == (B, K_MAX + 1) + + +def test_wgan_generator_gradients_flow(): + B = 4 + model = _small_generator() + cond_cont, cond_cat = _cond(B) + z = torch.randn(B, model.noise_dim) + gen_loss = model(z, cond_cont, cond_cat).sum() + nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum() + (gen_loss + nsec_loss).backward() + for name, p in model.named_parameters(): + assert p.grad is not None, f"no grad for {name}" + + +def test_critic_output_shape(): + B = 8 + critic = _small_critic() + cond_cont, cond_cat = _cond(B) + x = torch.randn(B, X_DIM) + out = critic(x, cond_cont, cond_cat) + assert out.shape == (B,) + + +def test_sample_wgan_shape(): + B = 6 + model = _small_generator() + cond_cont, cond_cat = _cond(B) + sample, n_sec = sample_wgan(model, cond_cont, cond_cat) + assert sample.shape == (B, X_DIM) + assert n_sec.shape == (B,) + + +# --- Stage-2 generator/critic --- + + +def test_wgan_secondary_generator_output_shape(): + B = 8 + model = _small_sec_generator() + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + z = torch.randn(B, model.noise_dim) + out = model(z, cond_cont, cond_cat, stage1_out) + assert out.shape == (B, SEC_DIM) + + +def test_secondary_critic_output_shape(): + B = 8 + critic = _small_sec_critic() + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + x = torch.randn(B, SEC_DIM) + out = critic(x, cond_cont, cond_cat, stage1_out) + assert out.shape == (B,) + + +def test_sample_secondaries_wgan_shape(): + B = 5 + model = _small_sec_generator() + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + n_sec_pred = torch.randint(0, K_MAX, (B,)) + sec_cont, sec_phys, sec_valid = sample_secondaries_wgan( + model, cond_cont, cond_cat, stage1_out, n_sec_pred + ) + assert sec_cont.shape == (B, K_MAX, 4) + assert sec_phys.shape == (B, K_MAX, 2) + assert sec_valid.shape == (B, K_MAX) + + +# --- Losses --- + + +def test_gradient_penalty_nonneg(): + B = 8 + critic = _small_critic() + cond_cont, cond_cat = _cond(B) + real = torch.randn(B, X_DIM) + fake = torch.randn(B, X_DIM) + gp = gradient_penalty(lambda x: critic(x, cond_cont, cond_cat), real, fake) + assert gp.item() >= 0.0 + assert gp.shape == () + + +def test_gradient_penalty_masked(): + B = 8 + sec_critic = _small_sec_critic() + cond_cont, cond_cat = _cond(B) + stage1_out = torch.randn(B, X_DIM) + n_sec = torch.randint(0, K_MAX, (B,)) + mask = _mask(B, n_sec) + real = torch.randn(B, SEC_DIM) * mask + fake = torch.randn(B, SEC_DIM) * mask + gp = gradient_penalty( + lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask + ) + assert gp.item() >= 0.0 + + +def test_critic_loss_scalar_and_grad(): + B = 8 + critic = _small_critic() + cond_cont, cond_cat = _cond(B) + real = torch.randn(B, X_DIM) + fake = torch.randn(B, X_DIM) + loss = critic_loss( + lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0 + ) + assert loss.shape == () + loss.backward() + assert any(p.grad is not None for p in critic.parameters()) + + +def test_generator_loss_scalar_and_grad(): + B = 4 + generator = _small_generator() + critic = _small_critic() + cond_cont, cond_cat = _cond(B) + z = torch.randn(B, generator.noise_dim) + fake = generator(z, cond_cont, cond_cat) + loss = generator_loss(lambda x: critic(x, cond_cont, cond_cat), fake) + assert loss.shape == () + loss.backward() + assert any(p.grad is not None for p in generator.parameters())