lars
87e37ebe14
CI / Format (ruff format) (push) Successful in 29s
CI / Lint (ruff check) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Lint (ruff check) (pull_request) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 4m32s
CI / Tests (pull_request) Successful in 4m25s
stage{1,2}_model.active = false already trains one stage alone, but the
checkpoint it writes holds only that stage, so giant rollout refuses it --
the "retrain stage 2 alone against a fixed, known-good stage 1" experiment
the 2026-08-03 species failure calls for wasn't runnable end to end.
Adds stage{1,2}_model.init_from (a checkpoint .pt to load this stage's
weights from before training) and .freeze (never update them), symmetric
across both stages. Both stages stay active = true, so both get built and
both land in the output checkpoint -- the frozen stage is merely
initialized from disk instead of from scratch.
Decisions made during planning:
- Soft freeze: forward/backward still run every batch (loss/grad_norm stay
meaningful, no autograd special-casing), only optimizer.step() (and, for
the frozen stage, lr_sched.step()/EMA update) is skipped -- weights are
byte-identical for the whole run. This is StageTrainer._step_optimizer,
shared by the non-adversarial path and both halves (generator + critic)
of the WGAN path, so a frozen WGAN stage's critic freezes too.
- validate_config requires init_from whenever freeze = true, unless the run
is a --resume (a resumed frozen stage's weights come from the resume
checkpoint instead) -- freezing a randomly-initialized model is almost
certainly a mistake.
- CLI flags on both `giant train` and `giant new-run`
(--stage{1,2}-init-from/--stage{1,2}-freeze), matching every other
per-stage model knob's existing treatment.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
525 lines
21 KiB
Python
525 lines
21 KiB
Python
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
from giant import config
|
|
from giant.constants import (
|
|
COND_DIM,
|
|
PARTICLE_PHYS_DIM,
|
|
X_DIM,
|
|
)
|
|
from giant.data import setup_cache
|
|
from giant.data.loader import (
|
|
TopNMap,
|
|
event_id_offset,
|
|
find_parquet_files,
|
|
iter_file_chunks,
|
|
build_index_maps_from_files,
|
|
build_pdg_topn_map_from_files,
|
|
build_process_map_from_files,
|
|
build_topn_map_from_files,
|
|
)
|
|
from giant.data.transforms import (
|
|
Normalizer,
|
|
build_features,
|
|
_WelfordAccumulator,
|
|
_ReservoirSampler,
|
|
sorted_membership,
|
|
)
|
|
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
|
from giant.model.network import build_models, build_critics, resolve_type_n_classes
|
|
from giant.training import train as run_training
|
|
|
|
|
|
@dataclass
|
|
class SetupStageResult:
|
|
"""Everything `run_train_job`'s pre-epoch setup stage derives from `data`.
|
|
|
|
Also returned standalone by `run_setup_stage` for callers (e.g. `dwarf
|
|
warm-cache`) that only want to populate/refresh the setup cache sidecar
|
|
without actually training a model.
|
|
"""
|
|
|
|
files: list[Path]
|
|
pdg_map: dict[int, int]
|
|
mat_map: dict[str, int]
|
|
proc_map: dict[str, int] | None
|
|
pdg_topn_map: TopNMap | None
|
|
sec_type_topn_map: TopNMap | None
|
|
mat_topn_map: TopNMap | None
|
|
cond_norm: Normalizer
|
|
tgt_norm: Normalizer
|
|
sec_phys_norm: Normalizer
|
|
train_events: set
|
|
val_events: set
|
|
n_train_steps: int
|
|
|
|
|
|
def _seed_energy_router(
|
|
router_cfg: dict,
|
|
cond_norm: Normalizer,
|
|
energy_quantiles: np.ndarray,
|
|
energy_idx: int,
|
|
echo,
|
|
) -> None:
|
|
"""Mutate `router_cfg["centers_init"]` in place from real data quantiles,
|
|
when this stage's router is an enabled EnergyRouter. Shared by both
|
|
stages' router configs — each seeded independently, since v0.3.0 stages
|
|
may have entirely different router configs."""
|
|
active = router_cfg.get("enabled") and router_cfg.get("type") == "energy"
|
|
if not active:
|
|
return
|
|
if energy_quantiles.size == 0:
|
|
echo(" warning: no energy samples collected — EnergyRouter falls back to default centers")
|
|
return
|
|
assert cond_norm.mean is not None and cond_norm.std is not None
|
|
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
|
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
|
|
centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
|
|
router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
|
|
echo(f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}")
|
|
|
|
|
|
def run_setup_stage(
|
|
data: str | Path,
|
|
val_fraction: float,
|
|
seed: int,
|
|
cfg: dict,
|
|
cache_setup: bool = True,
|
|
rebuild_setup_cache: bool = False,
|
|
echo=print,
|
|
) -> SetupStageResult:
|
|
"""Scan `data` for everything training needs before the epoch loop: the
|
|
train/val event split, pdg/material vocab maps, an optional process map
|
|
(needed if either stage's router is type="process"), and the Stage-1/
|
|
Stage-2 normalizers.
|
|
|
|
`cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/
|
|
`stage2_model`), already passed through `giant.config.validate_config`.
|
|
`conditioning.particle.type` and `conditioning.material.type` are
|
|
independent and may differ.
|
|
|
|
Reads from and writes to the `giant.data.setup_cache` sidecar when
|
|
`cache_setup` is set (`rebuild_setup_cache` ignores — but still
|
|
refreshes — any existing sidecar content). Each stage's `router` config
|
|
is mutated in place: an active `EnergyRouter` (`router.type == "energy"`)
|
|
gets its `centers_init` seeded from real data quantiles here.
|
|
"""
|
|
particle_conditioning = cfg["conditioning"]["particle"]["type"]
|
|
material_conditioning = cfg["conditioning"]["material"]["type"]
|
|
k_max = cfg["stage2_model"]["k_max"]
|
|
stage1_router = cfg["stage1_model"].get("router") or {}
|
|
stage2_router = cfg["stage2_model"].get("router") or {}
|
|
|
|
files = find_parquet_files(data)
|
|
echo(f"found {len(files)} parquet file(s)")
|
|
|
|
cache: setup_cache.SetupCache | None = None
|
|
if cache_setup:
|
|
if rebuild_setup_cache:
|
|
echo("setup cache: --rebuild-setup-cache given, recomputing all sections")
|
|
loaded = None
|
|
else:
|
|
loaded = setup_cache.load(data, files, echo=echo)
|
|
cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
|
|
|
|
if cache is not None and cache.event_index is not None:
|
|
unique_ids, counts = cache.event_index
|
|
echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
|
|
else:
|
|
echo("scanning event IDs …")
|
|
unique_ids, counts = setup_cache.compute_event_index_from_files(files)
|
|
if cache is not None:
|
|
cache.event_index = (unique_ids, counts)
|
|
|
|
train_events, val_events = make_event_split(unique_ids, val_fraction=val_fraction, seed=seed)
|
|
events_arr = np.array(sorted(train_events))
|
|
n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
|
|
echo(f" {int(counts.sum()):,} steps | {len(train_events)} train events | {len(val_events)} val events")
|
|
|
|
if cache is not None and cache.vocab is not None:
|
|
pdg_map, mat_map = cache.vocab
|
|
echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
|
|
else:
|
|
echo("building vocabulary maps …")
|
|
pdg_map, mat_map = build_index_maps_from_files(files)
|
|
echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
|
|
if cache is not None:
|
|
cache.vocab = (pdg_map, mat_map)
|
|
|
|
# A process map is needed if either stage's router reads the physics
|
|
# process label (type="process"). Only one map is built even if both
|
|
# stages want one — see the module-level note in giant/cli.py's
|
|
# _router_total_experts for why composed-router n_experts isn't a plain
|
|
# int; process routers are never composed in practice, so this doesn't
|
|
# need that generality.
|
|
proc_map: dict[str, int] | None = None
|
|
process_router_cfg = next(
|
|
(r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
|
|
None,
|
|
)
|
|
if process_router_cfg is not None:
|
|
n_experts = process_router_cfg["n_experts"]
|
|
cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
|
|
if cached_proc_map is not None:
|
|
proc_map = cached_proc_map
|
|
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
|
|
else:
|
|
echo("building process vocabulary …")
|
|
proc_map = build_process_map_from_files(files, n_experts=n_experts)
|
|
echo(f" {len(proc_map)} process labels mapped to {n_experts} experts")
|
|
if cache is not None:
|
|
cache.proc_maps[n_experts] = proc_map
|
|
|
|
# Top-N-plus-other maps for onehot conditioning/type axes.
|
|
# The PDG axis is used independently by conditioning.particle.type="onehot"
|
|
# (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
|
|
# (secondary-species decode) — their class counts can now differ (gitea
|
|
# #29: stage2_model.particle_type.n_classes, 0 = inherit
|
|
# conditioning.particle.emb_dim), so each is resolved and built
|
|
# independently via _pdg_topn below. cache.topn_maps is keyed by
|
|
# (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
|
|
# the same N the second call is a cache hit against the first — no extra
|
|
# scan in the common case where they still match. The material axis is
|
|
# independent of both.
|
|
particle_cfg = cfg["conditioning"]["particle"]
|
|
material_cfg = cfg["conditioning"]["material"]
|
|
particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
|
|
particle_type_target = particle_type_cfg.target
|
|
|
|
def _pdg_topn(n_classes: int) -> TopNMap:
|
|
cache_key = setup_cache.topn_key("pdg", n_classes)
|
|
cached = cache.topn_maps.get(cache_key) if cache is not None else None
|
|
if cached is not None:
|
|
echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
|
|
return cached
|
|
echo("building pdg top-N map …")
|
|
topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
|
|
echo(f" {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
|
|
if cache is not None:
|
|
cache.topn_maps[cache_key] = topn_map
|
|
return topn_map
|
|
|
|
pdg_topn_map: TopNMap | None = None
|
|
if particle_cfg["type"] == "onehot":
|
|
pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
|
|
|
|
sec_type_topn_map: TopNMap | None = None
|
|
if particle_type_target == "onehot":
|
|
sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
|
|
sec_type_topn_map = _pdg_topn(sec_type_n_classes)
|
|
|
|
mat_topn_map: TopNMap | None = None
|
|
if material_cfg["type"] == "onehot":
|
|
n_classes = material_cfg["emb_dim"]
|
|
cache_key = setup_cache.topn_key("material", n_classes)
|
|
cached = cache.topn_maps.get(cache_key) if cache is not None else None
|
|
if cached is not None:
|
|
mat_topn_map = cached
|
|
echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
|
|
else:
|
|
echo("building material top-N map …")
|
|
mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
|
|
echo(f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
|
|
if cache is not None:
|
|
cache.topn_maps[cache_key] = mat_topn_map
|
|
|
|
energy_router_active = any(r.get("enabled") and r.get("type") == "energy" for r in (stage1_router, stage2_router))
|
|
energy_idx = 3
|
|
norm_key = setup_cache.normalizer_key(val_fraction, seed, particle_conditioning, material_conditioning)
|
|
entry = cache.normalizers.get(norm_key) if cache is not None else None
|
|
|
|
if entry is not None:
|
|
echo(f"normalizer: cache hit (key={norm_key!r})")
|
|
cond_norm = entry.cond_norm
|
|
tgt_norm = entry.tgt_norm
|
|
sec_phys_norm = entry.sec_phys_norm
|
|
energy_quantiles = entry.energy_quantiles
|
|
else:
|
|
echo("fitting normalizer (streaming) …")
|
|
cond_acc = _WelfordAccumulator(COND_DIM)
|
|
tgt_acc = _WelfordAccumulator(X_DIM)
|
|
sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
|
|
# EnergyRouter's default center spread (linspace over [-2, 2]) assumes
|
|
# the z-normalized energy column is roughly uniform, which real energy
|
|
# spectra rarely are — collect a reservoir sample here (reusing this
|
|
# same pass, not a second scan), then collapse it to a fixed quantile
|
|
# grid (setup_cache.energy_quantiles_from_sample) so centers can
|
|
# instead be seeded from actual data quantiles below. Collected
|
|
# whenever the setup cache is being populated, not only when *this*
|
|
# run's router is energy-typed, so a later run enabling an energy
|
|
# router against this same (val_fraction, seed, conditioning) key
|
|
# never needs to rescan just to seed centers.
|
|
collect_energy_sample = energy_router_active or cache is not None
|
|
energy_sampler = _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
|
|
for i, path in enumerate(files):
|
|
for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
|
|
mask = sorted_membership(chunk["event_id"], events_arr)
|
|
if not mask.any():
|
|
continue
|
|
chunk_tr = {k: v[mask] for k, v in chunk.items()}
|
|
feats = build_features(
|
|
chunk_tr,
|
|
pdg_map,
|
|
mat_map,
|
|
proc_map=proc_map,
|
|
require_secondaries=True,
|
|
particle_conditioning=particle_conditioning,
|
|
material_conditioning=material_conditioning,
|
|
sec_phys_only=True,
|
|
# This pass reads only cond_cont/sec_cont, never cond_cat —
|
|
# but cond_cat's width is the conditioning modes' call
|
|
# (giant.cond_layout.CondLayout), so an "onehot" axis still
|
|
# has to be handed its map rather than silently yielding a
|
|
# narrower array.
|
|
pdg_topn_map=pdg_topn_map.class_map if pdg_topn_map is not None else None,
|
|
mat_topn_map=mat_topn_map.class_map if mat_topn_map is not None else None,
|
|
k_max=k_max,
|
|
)
|
|
cond_cont = feats.cond_cont
|
|
target_s1 = feats.target_s1
|
|
n_sec = feats.n_sec
|
|
sec_cont = feats.sec_cont
|
|
cond_acc.update(cond_cont)
|
|
tgt_acc.update(target_s1)
|
|
if energy_sampler is not None:
|
|
energy_sampler.update(cond_cont[:, energy_idx])
|
|
sec_valid = np.arange(sec_cont.shape[1])[None, :] < n_sec[:, None]
|
|
sec_phys = sec_cont[:, :, 4:6][sec_valid]
|
|
if len(sec_phys) > 0:
|
|
sec_phys_acc.update(sec_phys)
|
|
cond_norm = cond_acc.to_normalizer()
|
|
tgt_norm = tgt_acc.to_normalizer()
|
|
sec_phys_norm = sec_phys_acc.to_normalizer()
|
|
energy_quantiles = (
|
|
setup_cache.energy_quantiles_from_sample(energy_sampler.sample)
|
|
if energy_sampler is not None
|
|
else np.empty(0, dtype=np.float32)
|
|
)
|
|
if cache is not None:
|
|
cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
|
|
cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles
|
|
)
|
|
|
|
_seed_energy_router(stage1_router, cond_norm, energy_quantiles, energy_idx, echo)
|
|
_seed_energy_router(stage2_router, cond_norm, energy_quantiles, energy_idx, echo)
|
|
|
|
if cache is not None:
|
|
setup_cache.save(data, files, cache, echo=echo)
|
|
|
|
return SetupStageResult(
|
|
files=files,
|
|
pdg_map=pdg_map,
|
|
mat_map=mat_map,
|
|
proc_map=proc_map,
|
|
pdg_topn_map=pdg_topn_map,
|
|
sec_type_topn_map=sec_type_topn_map,
|
|
mat_topn_map=mat_topn_map,
|
|
cond_norm=cond_norm,
|
|
tgt_norm=tgt_norm,
|
|
sec_phys_norm=sec_phys_norm,
|
|
train_events=train_events,
|
|
val_events=val_events,
|
|
n_train_steps=n_train_steps,
|
|
)
|
|
|
|
|
|
def run_train_job(
|
|
data: Path,
|
|
cfg: dict,
|
|
out_dir: Path,
|
|
device: torch.device,
|
|
shuffle_buffer: int,
|
|
num_workers: int,
|
|
resume: Path | None = None,
|
|
cache_setup: bool = True,
|
|
rebuild_setup_cache: bool = False,
|
|
echo=print,
|
|
) -> None:
|
|
t = cfg["train"]
|
|
config.seed_everything(t["seed"])
|
|
|
|
out_dir = Path(out_dir)
|
|
|
|
# Soft warning (never blocks) — CLAUDE.md's Compute environment section
|
|
# asks that shared portal machines (portal1/deepthought{,2}/bms{1..3})
|
|
# stay within ~1/4 of CPU/RAM so as not to disturb other users' jobs;
|
|
# DataLoader's num_workers has no awareness of that on its own.
|
|
cpu_count = os.cpu_count() or 1
|
|
quota = max(1, cpu_count // 4)
|
|
if num_workers > quota:
|
|
echo(
|
|
f"warning: --num-workers={num_workers} exceeds ~1/4 of this "
|
|
f"machine's {cpu_count} CPU(s) ({quota}) — portal machines are "
|
|
"shared with other users (see CLAUDE.md's Compute environment "
|
|
"section)"
|
|
)
|
|
|
|
config.validate_config(cfg, resume=resume is not None)
|
|
particle_conditioning = cfg["conditioning"]["particle"]["type"]
|
|
material_conditioning = cfg["conditioning"]["material"]["type"]
|
|
k_max = cfg["stage2_model"]["k_max"]
|
|
setup = run_setup_stage(
|
|
data,
|
|
val_fraction=t["val_fraction"],
|
|
seed=t["seed"],
|
|
cfg=cfg,
|
|
cache_setup=cache_setup,
|
|
rebuild_setup_cache=rebuild_setup_cache,
|
|
echo=echo,
|
|
)
|
|
files = setup.files
|
|
pdg_map, mat_map, proc_map = setup.pdg_map, setup.mat_map, setup.proc_map
|
|
cond_norm, tgt_norm, sec_phys_norm = (
|
|
setup.cond_norm,
|
|
setup.tgt_norm,
|
|
setup.sec_phys_norm,
|
|
)
|
|
train_events, val_events, n_train_steps = (
|
|
setup.train_events,
|
|
setup.val_events,
|
|
setup.n_train_steps,
|
|
)
|
|
|
|
# cond_cat's onehot columns are present per-axis, independently, under
|
|
# that axis's own conditioning.{particle,material}.type == "onehot"
|
|
# (the two axes may mix freely). run_setup_stage builds each map
|
|
# whenever its own axis is "onehot" (see its own
|
|
# particle_cfg["type"]/material_cfg["type"]
|
|
# checks), so they're guaranteed non-None here — asserted, not just
|
|
# assumed, so a future wiring bug fails loudly instead of silently
|
|
# dropping the onehot columns.
|
|
cond_pdg_topn = None
|
|
cond_mat_topn = None
|
|
if particle_conditioning == "onehot":
|
|
assert setup.pdg_topn_map is not None
|
|
cond_pdg_topn = setup.pdg_topn_map.class_map
|
|
if material_conditioning == "onehot":
|
|
assert setup.mat_topn_map is not None
|
|
cond_mat_topn = setup.mat_topn_map.class_map
|
|
|
|
# The secondary type-index map depends on stage2_model.particle_type.target,
|
|
# independently of conditioning's own onehot/embedding choice above
|
|
# (physical stays untouched/None).
|
|
particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target
|
|
if particle_type_target == "onehot":
|
|
assert setup.sec_type_topn_map is not None
|
|
sec_type_class_map = setup.sec_type_topn_map.class_map
|
|
elif particle_type_target == "embedding":
|
|
sec_type_class_map = pdg_map
|
|
else:
|
|
sec_type_class_map = None
|
|
|
|
total_train_batches = n_train_steps // t["batch_size"]
|
|
echo(f" ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")
|
|
|
|
train_ds = StreamingStepsDataset(
|
|
files=files,
|
|
split_events=train_events,
|
|
pdg_map=pdg_map,
|
|
mat_map=mat_map,
|
|
cond_normalizer=cond_norm,
|
|
target_normalizer=tgt_norm,
|
|
batch_size=t["batch_size"],
|
|
shuffle_buffer=shuffle_buffer,
|
|
shuffle=True,
|
|
proc_map=proc_map,
|
|
particle_conditioning=particle_conditioning,
|
|
material_conditioning=material_conditioning,
|
|
sec_phys_normalizer=sec_phys_norm,
|
|
pdg_topn_map=cond_pdg_topn,
|
|
mat_topn_map=cond_mat_topn,
|
|
sec_type_class_map=sec_type_class_map,
|
|
k_max=k_max,
|
|
)
|
|
val_ds = StreamingStepsDataset(
|
|
files=files,
|
|
split_events=val_events,
|
|
pdg_map=pdg_map,
|
|
mat_map=mat_map,
|
|
cond_normalizer=cond_norm,
|
|
target_normalizer=tgt_norm,
|
|
batch_size=t["batch_size"],
|
|
shuffle=False,
|
|
proc_map=proc_map,
|
|
particle_conditioning=particle_conditioning,
|
|
material_conditioning=material_conditioning,
|
|
sec_phys_normalizer=sec_phys_norm,
|
|
pdg_topn_map=cond_pdg_topn,
|
|
mat_topn_map=cond_mat_topn,
|
|
sec_type_class_map=sec_type_class_map,
|
|
k_max=k_max,
|
|
)
|
|
|
|
pin = device.type == "cuda"
|
|
train_loader = DataLoader(
|
|
train_ds,
|
|
batch_size=None,
|
|
num_workers=num_workers,
|
|
pin_memory=pin,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_ds,
|
|
batch_size=None,
|
|
num_workers=num_workers,
|
|
pin_memory=pin,
|
|
)
|
|
|
|
model_config = {
|
|
"pdg_vocab": len(pdg_map),
|
|
"mat_vocab": len(mat_map),
|
|
"conditioning": cfg["conditioning"],
|
|
"stage1_model": cfg["stage1_model"],
|
|
"stage2_model": cfg["stage2_model"],
|
|
}
|
|
|
|
models = build_models(model_config)
|
|
critics = build_critics(model_config)
|
|
for name, model in models.items():
|
|
if model is not None:
|
|
echo(f"{name}: {sum(p.numel() for p in model.parameters()):,} parameters")
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
meta = config.build_run_meta(
|
|
data=data,
|
|
seed=t["seed"],
|
|
n_pdg_codes=len(pdg_map),
|
|
n_materials=len(mat_map),
|
|
n_train_events=len(train_events),
|
|
n_val_events=len(val_events),
|
|
n_train_steps=n_train_steps,
|
|
)
|
|
config.save_config(cfg, out_dir, meta)
|
|
|
|
run_training(
|
|
cfg=cfg,
|
|
models=models,
|
|
critics=critics,
|
|
train_loader=train_loader,
|
|
val_loader=val_loader,
|
|
device=device,
|
|
out_dir=out_dir,
|
|
normalizer_dict={
|
|
"cond": cond_norm.to_dict(),
|
|
"target": tgt_norm.to_dict(),
|
|
"sec_phys": sec_phys_norm.to_dict(),
|
|
},
|
|
pdg_map={str(k): v for k, v in pdg_map.items()},
|
|
mat_map={str(k): v for k, v in mat_map.items()},
|
|
proc_map=proc_map,
|
|
pdg_topn_map=setup.pdg_topn_map,
|
|
sec_type_topn_map=setup.sec_type_topn_map,
|
|
mat_topn_map=setup.mat_topn_map,
|
|
model_config=model_config,
|
|
resume_path=resume,
|
|
total_train_batches=total_train_batches,
|
|
use_wandb=t.get("wandb", True),
|
|
wandb_project=t.get("wandb_project", "giant"),
|
|
wandb_run_name=t.get("wandb_run_name", ""),
|
|
wandb_log_every=t.get("wandb_log_every", 50),
|
|
)
|