diff --git a/giant/model/builders.py b/giant/model/builders.py index 56c20bc..0849db2 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -202,6 +202,8 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: cond_out_dim=cond_out_dim, dropout=s1_spec.dropout, stage="stage1", + trunk_type=s1_spec.trunk.type, + block_conditioning=s1_spec.trunk.block_conditioning, ) if s2_spec.active and build_objective(s2_spec.generator).is_adversarial: @@ -225,6 +227,8 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: dropout=s2_spec.dropout, stage="stage2", context_dim=s2_spec.context_dim, + trunk_type=s2_spec.trunk.type, + block_conditioning=s2_spec.trunk.block_conditioning, ) return result diff --git a/giant/model/models.py b/giant/model/models.py index 4284e5e..8fb91c3 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -8,7 +8,7 @@ from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM from giant.model.encoders import ConditionEncoder from giant.model.history import HistoryEncoder, build_history -from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head +from giant.model.layers import ContextAdapter, SinusoidalEmbedding, build_mlp_head from giant.model.objectives import build_objective from giant.model.routers import Router from giant.model.trunks import build_trunk @@ -188,6 +188,26 @@ class StageModel(nn.Module): hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio)) self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth) + def _build_context_fusion(self, x_dim: int, context_dim: int, cond_out_dim: int) -> None: + """Builds `self.context_adapter`/`self.fuse` — the stage-2-style + context-fusion pattern (project the previous stage's outcome down to + `context_dim` via `ContextAdapter`, concat onto the base conditioning, + project back to `cond_out_dim`) shared by `Stage2OneShot` and a + `stage="stage2"` `CriticModel` (gitea #57). Call from a subclass's + `__init__` before using `_cond_embed`.""" + self.context_adapter = ContextAdapter(x_dim, context_dim) + self.fuse = nn.Sequential( + nn.Linear(cond_out_dim + context_dim, cond_out_dim), + nn.SiLU(), + ) + + def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: + """Fuses base conditioning with the previous stage's outcome — pairs + with `_build_context_fusion`.""" + base = self.cond_enc(cond_cont, cond_cat) + ctx = self.context_adapter(stage1_out) + return self.fuse(torch.cat([base, ctx], dim=-1)) + def _require_n_sec_head(self) -> None: if self.n_sec_head is None: raise RuntimeError( @@ -363,11 +383,7 @@ class Stage2OneShot(StageModel): particle_type_cfg=particle_type_cfg, cond_enc=cond_enc, ) - self.context_adapter = ContextAdapter(x_dim, context_dim) - self.fuse = nn.Sequential( - nn.Linear(cond_out_dim + context_dim, cond_out_dim), - nn.SiLU(), - ) + self._build_context_fusion(x_dim, context_dim, cond_out_dim) target = self.particle_type_cfg.target type_head_out_dim = None if target == "physical" else k_max * self.type_dim self._build_trunk_and_heads( @@ -386,11 +402,6 @@ class Stage2OneShot(StageModel): type_head_cfg=type_head_cfg, ) - def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: - base = self.cond_enc(cond_cont, cond_cat) - ctx = self.context_adapter(stage1_out) - return self.fuse(torch.cat([base, ctx], dim=-1)) - def forward( self, x_t: torch.Tensor, @@ -702,11 +713,25 @@ class Stage2Autoregressive(StageModel): return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K) -class CriticModel(nn.Module): +class CriticModel(StageModel): """Generator-agnostic WGAN-GP critic body: a scalar realism score, for either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"` mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as - `Stage2OneShot`). Used only when that stage's `generator == "wgan"`.""" + `Stage2OneShot`, via `StageModel._build_context_fusion`/`_cond_embed`). + Used only when that stage's `generator == "wgan"`. + + Subclasses `StageModel` for the `cond_enc` construction and (stage 2) + context-fusion scaffolding only — its trunk is built directly via + `build_trunk` (output width 1) rather than through + `_build_trunk_and_heads`, since that helper is shaped around a + generator's `Objective`/time-embedding/flow-matching concerns + (`forward`'s `(x_t, cond) -> vector` shape) that don't apply to a critic's + `(x, cond) -> scalar` (gitea #57). `generator="wgan"` is passed to the + base purely because that's factually when a critic exists; nothing here + ever calls `_build_trunk_and_heads`, so no head/time-embedding machinery + is built from it. Never routed (MoE) — that's a separate, unrequested + axis of scope; see gitea #57's proposal, which covers only the trunk/ + block registries.""" def __init__( self, @@ -722,22 +747,26 @@ class CriticModel(nn.Module): stage: str = "stage1", context_dim: int = 64, context_in_dim: int = X_DIM, + trunk_type: str = "resmlp", + block_conditioning: str = "add", ) -> None: - super().__init__() + super().__init__( + pdg_vocab, + mat_vocab, + particle_cfg, + material_cfg, + cond_out_dim=cond_out_dim, + generator="wgan", + noise_dim=0, + ) if stage not in ("stage1", "stage2"): raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}") self.stage = stage - self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) if stage == "stage2": - self.context_adapter = ContextAdapter(context_in_dim, context_dim) - self.fuse = nn.Sequential( - nn.Linear(cond_out_dim + context_dim, cond_out_dim), - nn.SiLU(), - ) - self.input_proj = nn.Linear(in_dim, hidden_dim) - self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)]) - self.out_norm = nn.LayerNorm(hidden_dim) - self.out_proj = nn.Linear(hidden_dim, 1) + self._build_context_fusion(context_in_dim, context_dim, cond_out_dim) + self.trunk = build_trunk( + None, trunk_type, in_dim, 1, hidden_dim, n_res_blocks, cond_out_dim, dropout, block_conditioning + ) def forward( self, @@ -746,13 +775,9 @@ class CriticModel(nn.Module): cond_cat: torch.Tensor, stage1_out: torch.Tensor | None = None, ) -> torch.Tensor: - base = self.cond_enc(cond_cont, cond_cat) if self.stage == "stage2": - ctx = self.context_adapter(stage1_out) - cond = self.fuse(torch.cat([base, ctx], dim=-1)) + assert stage1_out is not None, "stage='stage2' CriticModel requires stage1_out" + cond = self._cond_embed(cond_cont, cond_cat, stage1_out) else: - cond = base - h = self.input_proj(x) - for block in self.blocks: - h = block(h, cond) - return self.out_proj(self.out_norm(h)).squeeze(-1) + cond = self.cond_enc(cond_cont, cond_cat) + return self.trunk(x, cond, cond_cont, cond_cat).squeeze(-1) diff --git a/tests/test_network.py b/tests/test_network.py index bcd1877..384a6b3 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -8,7 +8,10 @@ from giant.model.network import ( HISTORY_REGISTRY, AttentionHistory, ConditionEncoder, + CriticModel, + FilmResBlock, HistoryEncoder, + LinearTrunk, MarkovHistory, NoHistory, SinusoidalEmbedding, @@ -944,7 +947,7 @@ def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim(): assert wider_critic is not None # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding — # widening n_classes alone (emb_dim stays 4) must widen the critic input. - assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features + assert wider_critic.trunk.input_proj.in_features > default_n_classes_critic.trunk.input_proj.in_features # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─ @@ -1021,12 +1024,12 @@ def test_build_critics_omitted_particle_type_matches_default_config(): cfg["stage2_model"]["generator"] = "wgan" onehot_critic = build_critics(cfg)["stage2"] assert onehot_critic is not None - onehot_in_dim = onehot_critic.input_proj.in_features + onehot_in_dim = onehot_critic.trunk.input_proj.in_features cfg["stage2_model"]["particle_type"] = {"target": "physical"} physical_critic = build_critics(cfg)["stage2"] assert physical_critic is not None - physical_in_dim = physical_critic.input_proj.in_features + physical_in_dim = physical_critic.trunk.input_proj.in_features # onehot's per-slot type width is emb_dim classes vs. physical's fixed # (log-mass, charge) pair — different unless emb_dim happens to be 2, so @@ -1046,15 +1049,15 @@ def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_genera inherited = build_critics(cfg)["stage1"] assert inherited is not None - assert inherited.input_proj.out_features == 8 - assert len(inherited.blocks) == 1 + assert inherited.trunk.input_proj.out_features == 8 + assert len(inherited.trunk.blocks) == 1 cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16 cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3 overridden = build_critics(cfg)["stage1"] assert overridden is not None - assert overridden.input_proj.out_features == 16 - assert len(overridden.blocks) == 3 + assert overridden.trunk.input_proj.out_features == 16 + assert len(overridden.trunk.blocks) == 3 def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size(): @@ -1065,15 +1068,15 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera inherited = build_critics(cfg)["stage2"] assert inherited is not None - assert inherited.input_proj.out_features == 8 - assert len(inherited.blocks) == 1 + assert inherited.trunk.input_proj.out_features == 8 + assert len(inherited.trunk.blocks) == 1 cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16 cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3 overridden = build_critics(cfg)["stage2"] assert overridden is not None - assert overridden.input_proj.out_features == 16 - assert len(overridden.blocks) == 3 + assert overridden.trunk.input_proj.out_features == 16 + assert len(overridden.trunk.blocks) == 3 # ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive @@ -1234,3 +1237,34 @@ def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator): assert model.generator_kind == generator assert model.noise_dim == 8 assert (model.time_emb is not None) == build_objective(generator).needs_time + + +# ── CriticModel uses the trunk/block registries + StageModel base (gitea #57) ─ + + +def test_critic_model_is_stagemodel_subclass(): + assert issubclass(CriticModel, StageModel) + + +@pytest.mark.parametrize("stage", ["stage1", "stage2"]) +def test_build_critics_threads_trunk_type_from_generator_config(stage): + cfg = _minimal_model_config(share_stages=False) + cfg["stage1_model"]["generator"] = "wgan" + cfg["stage2_model"]["generator"] = "wgan" + cfg[f"{stage}_model"]["trunk"] = {"type": "linear"} + + critic = build_critics(cfg)[stage] + assert critic is not None + assert isinstance(critic.trunk, LinearTrunk) + + +@pytest.mark.parametrize("stage", ["stage1", "stage2"]) +def test_build_critics_threads_block_conditioning_from_generator_config(stage): + cfg = _minimal_model_config(share_stages=False) + cfg["stage1_model"]["generator"] = "wgan" + cfg["stage2_model"]["generator"] = "wgan" + cfg[f"{stage}_model"]["trunk"] = {"block_conditioning": "film"} + + critic = build_critics(cfg)[stage] + assert critic is not None + assert all(isinstance(block, FilmResBlock) for block in critic.trunk.blocks) diff --git a/tests/test_wgan.py b/tests/test_wgan.py index dcb2cb3..c911b7a 100644 --- a/tests/test_wgan.py +++ b/tests/test_wgan.py @@ -2,7 +2,7 @@ import torch from giant.config import ConditioningAxisConfig from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM -from giant.model.network import CriticModel, Stage1Model, Stage2OneShot +from giant.model.network import CriticModel, LinearTrunk, Stage1Model, Stage2OneShot from giant.model.wgan import critic_loss, generator_loss, gradient_penalty from giant.sample import sample_secondaries_wgan, sample_wgan @@ -115,6 +115,33 @@ def test_critic_output_shape(): assert out.shape == (B,) +def test_critic_model_honours_trunk_type_and_block_conditioning(): + """gitea #57: CriticModel routes its body through build_trunk/build_block + like every generator stage model, instead of hand-rolling a plain + ResBlock stack.""" + B = 8 + critic = CriticModel( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + in_dim=X_DIM, + hidden_dim=32, + n_res_blocks=2, + stage="stage1", + trunk_type="linear", + block_conditioning="adaln", + ) + assert isinstance(critic.trunk, LinearTrunk) + 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) + loss.backward() + for name, p in critic.named_parameters(): + assert p.grad is not None, f"no grad for {name}" + + def test_sample_wgan_shape(): B = 6 model = _small_generator()