lars
4fc15ecdfc
CI / Lint (ruff check) (push) Successful in 26s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 28s
CI / Type check (ty) (push) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 35s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 36s
CI / Tests (pull_request) Successful in 1m41s
CI / Tests (push) Successful in 1m47s
Builds the shared top-N-plus-other PDG/material maps (pooling both primary
and secondary occurrences for PDG, directly targeting the meeting's
species-collapse failure mode) and wires up conditioning.{particle,material}
= "onehot" plus stage2_model.particle_type.target in ("onehot", "embedding")
end-to-end: setup-cache persistence, Stage2OneShot's type_head (flow/ddpm)
vs. folded+ST-Gumbel-relaxed adversarial slice (wgan), and the corresponding
CE/MSE training losses. particle_type.target = "physical" stays byte-for-byte
unchanged, keeping the v0.2 migration shim's bit-identical guarantee intact.
giant predict/rollout fail loudly on a onehot/embedding checkpoint until
full decode support lands in step 6.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
403 lines
13 KiB
Python
403 lines
13 KiB
Python
"""Tests for giant/train.py."""
|
|
|
|
import copy
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
|
|
from giant.model.network import build_critics, build_models
|
|
from giant.train import (
|
|
FlowDDPMStageTrainer,
|
|
WGANStageTrainer,
|
|
_gumbel_tau,
|
|
_wandb_run_config,
|
|
train,
|
|
)
|
|
|
|
PDG_VOCAB = 6
|
|
MAT_VOCAB = 3
|
|
|
|
|
|
def test_gumbel_tau_at_step_zero_is_start():
|
|
assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0
|
|
|
|
|
|
def test_gumbel_tau_at_total_steps_is_end():
|
|
assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9
|
|
|
|
|
|
def test_gumbel_tau_interpolates_linearly_midway():
|
|
assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9
|
|
|
|
|
|
def test_gumbel_tau_clamps_beyond_total_steps():
|
|
assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)
|
|
|
|
|
|
def test_gumbel_tau_handles_zero_total_steps():
|
|
# total_steps=0 is guarded to 1 internally: step=0 gives zero progress
|
|
# (still tau_start), any step>=1 immediately clamps to full progress.
|
|
assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
|
|
assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9
|
|
|
|
|
|
def test_wandb_run_config_includes_full_cfg_and_param_counts():
|
|
cfg = {
|
|
"train": {"lr": 3e-4},
|
|
"conditioning": {"out_dim": 128},
|
|
"stage1_model": {"generator": "flow"},
|
|
"stage2_model": {"generator": "wgan"},
|
|
}
|
|
wcfg = _wandb_run_config(
|
|
cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}
|
|
)
|
|
assert wcfg["train"] == {"lr": 3e-4}
|
|
assert wcfg["stage1_model"] == {"generator": "flow"}
|
|
assert wcfg["stage2_model"] == {"generator": "wgan"}
|
|
assert wcfg["model_config"] == {"pdg_vocab": 3}
|
|
assert wcfg["param_counts"] == {"stage1": 100}
|
|
|
|
|
|
def test_wandb_run_config_handles_missing_model_config():
|
|
cfg = {"train": {}, "conditioning": {}, "stage1_model": {}, "stage2_model": {}}
|
|
wcfg = _wandb_run_config(cfg, model_config=None, param_counts={})
|
|
assert wcfg["model_config"] == {}
|
|
|
|
|
|
# --- end-to-end train() integration tests -----------------------------------
|
|
|
|
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
|
|
|
|
|
def _base_cfg():
|
|
return {
|
|
"conditioning": {
|
|
"out_dim": 32,
|
|
"share_stages": False,
|
|
"particle": dict(PARTICLE_CFG),
|
|
"material": dict(MATERIAL_CFG),
|
|
},
|
|
"stage1_model": {
|
|
"active": True,
|
|
"generator": "flow",
|
|
"hidden_dim": 24,
|
|
"n_res_blocks": 2,
|
|
"dropout": 0.0,
|
|
"lambda": 1.0,
|
|
"flow": {"time_dim": 16},
|
|
"ddpm": {"time_dim": 16, "n_steps": 50},
|
|
"wgan": {
|
|
"noise_dim": 16,
|
|
"n_critic": 2,
|
|
"gp_weight": 10.0,
|
|
"critic_lr": 0.0,
|
|
},
|
|
"router": {"enabled": False},
|
|
},
|
|
"stage2_model": {
|
|
"active": True,
|
|
"decoder": "one_shot",
|
|
"generator": "wgan",
|
|
"hidden_dim": 24,
|
|
"n_res_blocks": 2,
|
|
"dropout": 0.0,
|
|
"lambda": 1.0,
|
|
"k_max": K_MAX,
|
|
"context_dim": 16,
|
|
"n_sec": {"mode": "head", "lambda": 0.1},
|
|
"flow": {"time_dim": 16},
|
|
"ddpm": {"time_dim": 16, "n_steps": 50},
|
|
"wgan": {
|
|
"noise_dim": 16,
|
|
"n_critic": 2,
|
|
"gp_weight": 10.0,
|
|
"critic_lr": 0.0,
|
|
},
|
|
"router": {"enabled": False, "tie_to_stage1": False},
|
|
},
|
|
"train": {
|
|
"epochs": 2,
|
|
"batch_size": 8,
|
|
"lr": 3e-4,
|
|
"weight_decay": 0.01,
|
|
"ema_decay": 0.999,
|
|
"warmup_epochs": 0,
|
|
"val_fraction": 0.1,
|
|
"max_val_batches": 0,
|
|
"num_workers": 0,
|
|
"seed": 0,
|
|
"validate_every": 0,
|
|
"validate_steps": 2,
|
|
"wandb": False,
|
|
},
|
|
}
|
|
|
|
|
|
def _fake_batches(n_batches, batch_size, seed=0):
|
|
g = torch.Generator().manual_seed(seed)
|
|
batches = []
|
|
for _ in range(n_batches):
|
|
cond_cont = torch.randn(batch_size, COND_DIM, generator=g)
|
|
cond_cat = torch.stack(
|
|
[
|
|
torch.randint(0, PDG_VOCAB, (batch_size,), generator=g),
|
|
torch.randint(0, MAT_VOCAB, (batch_size,), generator=g),
|
|
],
|
|
dim=1,
|
|
)
|
|
x1 = torch.randn(batch_size, X_DIM, generator=g)
|
|
n_sec = torch.randint(0, K_MAX, (batch_size,), generator=g)
|
|
sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
|
|
proc_idx = torch.zeros(batch_size, dtype=torch.long)
|
|
sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
|
|
batches.append(
|
|
(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
|
|
)
|
|
return batches
|
|
|
|
|
|
def _model_config(cfg):
|
|
return {
|
|
"pdg_vocab": PDG_VOCAB,
|
|
"mat_vocab": MAT_VOCAB,
|
|
"conditioning": cfg["conditioning"],
|
|
"stage1_model": cfg["stage1_model"],
|
|
"stage2_model": cfg["stage2_model"],
|
|
}
|
|
|
|
|
|
def _run_train(cfg, out_dir, resume_path=None):
|
|
model_config = _model_config(cfg)
|
|
models = build_models(model_config)
|
|
critics = build_critics(model_config)
|
|
train_loader = _fake_batches(4, cfg["train"]["batch_size"])
|
|
val_loader = _fake_batches(2, cfg["train"]["batch_size"], seed=1)
|
|
train(
|
|
cfg=cfg,
|
|
models=models,
|
|
critics=critics,
|
|
train_loader=train_loader,
|
|
val_loader=val_loader,
|
|
device=torch.device("cpu"),
|
|
out_dir=out_dir,
|
|
normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}},
|
|
pdg_map={"22": 0},
|
|
mat_map={"G4_AIR": 0},
|
|
proc_map=None,
|
|
model_config=model_config,
|
|
total_train_batches=4,
|
|
resume_path=resume_path,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"label,mutate",
|
|
[
|
|
("both_flow", lambda cfg: None),
|
|
("both_wgan", lambda cfg: cfg["stage1_model"].__setitem__("generator", "wgan")),
|
|
(
|
|
"mixed_stage1_flow_stage2_wgan",
|
|
lambda cfg: None, # already the default
|
|
),
|
|
(
|
|
"mixed_stage1_wgan_stage2_flow",
|
|
lambda cfg: (
|
|
cfg["stage1_model"].__setitem__("generator", "wgan"),
|
|
cfg["stage2_model"].__setitem__("generator", "flow"),
|
|
),
|
|
),
|
|
("stage1_only", lambda cfg: cfg["stage2_model"].__setitem__("active", False)),
|
|
("stage2_only", lambda cfg: cfg["stage1_model"].__setitem__("active", False)),
|
|
(
|
|
"both_ddpm_stage1_flow_stage2",
|
|
lambda cfg: (
|
|
cfg["stage1_model"].__setitem__("generator", "ddpm"),
|
|
cfg["stage2_model"].__setitem__("generator", "flow"),
|
|
),
|
|
),
|
|
(
|
|
"routed_stage1_energy_gumbel",
|
|
lambda cfg: cfg["stage1_model"].__setitem__(
|
|
"router",
|
|
{
|
|
"enabled": True,
|
|
"type": "energy",
|
|
"n_experts": 3,
|
|
"temperature": 0.5,
|
|
"learn_centers": True,
|
|
"lambda_balance": 0.1,
|
|
"lambda_entropy": 0.01,
|
|
"gumbel": True,
|
|
"gumbel_tau_start": 1.0,
|
|
"gumbel_tau_end": 0.1,
|
|
},
|
|
),
|
|
),
|
|
(
|
|
"stage2_onehot_target_wgan",
|
|
lambda cfg: cfg["stage2_model"].__setitem__(
|
|
"particle_type", {"target": "onehot", "lambda": 1.0}
|
|
),
|
|
),
|
|
(
|
|
"stage2_onehot_target_flow",
|
|
lambda cfg: (
|
|
cfg["stage2_model"].__setitem__("generator", "flow"),
|
|
cfg["stage2_model"].__setitem__(
|
|
"particle_type", {"target": "onehot", "lambda": 1.0}
|
|
),
|
|
),
|
|
),
|
|
(
|
|
"stage2_embedding_target_wgan",
|
|
lambda cfg: (
|
|
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
|
|
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
|
|
cfg["stage2_model"].__setitem__(
|
|
"particle_type", {"target": "embedding", "lambda": 1.0}
|
|
),
|
|
),
|
|
),
|
|
(
|
|
"stage2_embedding_target_flow",
|
|
lambda cfg: (
|
|
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
|
|
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
|
|
cfg["stage2_model"].__setitem__("generator", "flow"),
|
|
cfg["stage2_model"].__setitem__(
|
|
"particle_type", {"target": "embedding", "lambda": 1.0}
|
|
),
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_train_end_to_end(label, mutate):
|
|
cfg = _base_cfg()
|
|
mutate(cfg)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
out_dir = Path(tmp) / "run"
|
|
_run_train(cfg, out_dir)
|
|
assert (out_dir / "last.pt").exists()
|
|
assert (out_dir / "metrics.csv").exists()
|
|
ckpt = torch.load(out_dir / "last.pt", weights_only=False)
|
|
if cfg["stage1_model"]["active"]:
|
|
assert "model" in ckpt
|
|
else:
|
|
assert "model" not in ckpt
|
|
if cfg["stage2_model"]["active"]:
|
|
assert "sec_decoder" in ckpt
|
|
else:
|
|
assert "sec_decoder" not in ckpt
|
|
|
|
|
|
def test_train_resume_continues_from_checkpoint():
|
|
cfg = _base_cfg()
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
out_dir = Path(tmp) / "run"
|
|
_run_train(cfg, out_dir)
|
|
ckpt_before = torch.load(out_dir / "last.pt", weights_only=False)
|
|
assert ckpt_before["epoch"] == 2
|
|
|
|
cfg2 = copy.deepcopy(cfg)
|
|
cfg2["train"]["epochs"] = 3
|
|
_run_train(cfg2, out_dir, resume_path=out_dir / "last.pt")
|
|
ckpt_after = torch.load(out_dir / "last.pt", weights_only=False)
|
|
assert ckpt_after["epoch"] == 3
|
|
assert ckpt_after["global_step"] > ckpt_before["global_step"]
|
|
|
|
|
|
def test_train_raises_when_no_active_stage():
|
|
cfg = _base_cfg()
|
|
cfg["stage1_model"]["active"] = False
|
|
cfg["stage2_model"]["active"] = False
|
|
model_config = _model_config(cfg)
|
|
models = build_models(model_config)
|
|
critics = build_critics(model_config)
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
with pytest.raises(ValueError, match="no active stage"):
|
|
train(
|
|
cfg=cfg,
|
|
models=models,
|
|
critics=critics,
|
|
train_loader=_fake_batches(1, 8),
|
|
val_loader=_fake_batches(1, 8),
|
|
device=torch.device("cpu"),
|
|
out_dir=Path(tmp) / "run",
|
|
model_config=model_config,
|
|
total_train_batches=1,
|
|
)
|
|
|
|
|
|
def test_metrics_csv_columns_are_stage_prefixed():
|
|
cfg = _base_cfg()
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
out_dir = Path(tmp) / "run"
|
|
_run_train(cfg, out_dir)
|
|
header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
|
|
assert "stage1_train_loss" in header
|
|
assert "stage2_train_d_loss" in header
|
|
assert "val_loss" in header
|
|
assert "epoch" in header
|
|
|
|
|
|
def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
|
|
"""Regression test: on a non-generator-step batch, if this stage's model
|
|
has no n_sec_head (n_sec defaults to stage 2, decision 1), g_loss is a
|
|
graph-less zero — .backward() must not be called on it."""
|
|
cfg = _base_cfg()
|
|
cfg["stage1_model"]["generator"] = "wgan"
|
|
model_config = _model_config(cfg)
|
|
models = build_models(model_config)
|
|
critics = build_critics(model_config)
|
|
assert models["stage1"] is not None and critics["stage1"] is not None
|
|
trainer = WGANStageTrainer(
|
|
name="stage1",
|
|
model=models["stage1"],
|
|
critic=critics["stage1"],
|
|
is_stage2=False,
|
|
lambda_weight=1.0,
|
|
n_sec_lambda=0.1,
|
|
n_critic=1000, # never a generator step in this test
|
|
gp_weight=10.0,
|
|
lr=3e-4,
|
|
critic_lr=0.0,
|
|
ema_decay=0.0,
|
|
warmup_epochs=0,
|
|
epochs=1,
|
|
steps_per_epoch=4,
|
|
device=torch.device("cpu"),
|
|
)
|
|
assert trainer.model.n_sec_head is None
|
|
batch = _fake_batches(1, 8)[0]
|
|
stats = trainer.step(batch, torch.device("cpu"), global_step=1)
|
|
assert stats["did_g_step"] is False
|
|
|
|
|
|
def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
|
|
with pytest.raises(NotImplementedError):
|
|
FlowDDPMStageTrainer(
|
|
name="stage2",
|
|
model=torch.nn.Linear(1, 1),
|
|
is_stage2=True,
|
|
generator="ddpm",
|
|
lambda_weight=1.0,
|
|
n_sec_lambda=0.1,
|
|
lambda_balance=0.0,
|
|
lambda_proc=0.0,
|
|
lambda_entropy=0.0,
|
|
gumbel_tau_start=1.0,
|
|
gumbel_tau_end=0.1,
|
|
lr=3e-4,
|
|
weight_decay=0.01,
|
|
ema_decay=0.0,
|
|
warmup_epochs=0,
|
|
epochs=1,
|
|
steps_per_epoch=1,
|
|
ddpm_n_steps=50,
|
|
device=torch.device("cpu"),
|
|
)
|