| |
| """Run one deterministic Adam step for legacy or rewritten ROOT-GNN.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import dgl |
| import numpy as np |
| import torch |
|
|
| from validation.artifacts import load_artifact |
| from validation.forward import _graph |
|
|
|
|
| def _weighted_loss(logits, labels, weights): |
| elementwise = torch.nn.functional.cross_entropy(logits, labels, reduction="none") |
| result = logits.new_zeros(()) |
| for label in torch.unique(labels): |
| mask = labels == label |
| result = ( |
| result + (weights[mask] * elementwise[mask]).sum() / weights[mask].sum() |
| ) |
| return result / len(torch.unique(labels)) |
|
|
|
|
| def _flat_state(model, *, gradients: bool = False): |
| values = [] |
| for name, parameter in sorted(model.named_parameters()): |
| if gradients: |
| if parameter.grad is None: |
| raise RuntimeError(f"missing gradient for {name}") |
| values.append(parameter.grad.detach().cpu().reshape(-1)) |
| else: |
| values.append(parameter.detach().cpu().reshape(-1)) |
| return torch.cat(values).numpy() |
|
|
|
|
| def _flat_optimizer_state(model, optimizer, key): |
| values = [] |
| steps = [] |
| for name, parameter in sorted(model.named_parameters()): |
| state = optimizer.state[parameter] |
| values.append(state[key].detach().cpu().reshape(-1)) |
| steps.append(float(state["step"])) |
| return torch.cat(values).numpy(), np.asarray(steps, dtype=np.float64) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--artifact", type=Path, required=True) |
| parser.add_argument("--checkpoint", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--events", type=int, default=8) |
| parser.add_argument("--steps", type=int, default=1) |
| args = parser.parse_args() |
| artifact = load_artifact(args.artifact) |
| count = min(args.events, artifact.event_count) |
| graphs = [_graph(artifact, index) for index in range(count)] |
| graph = dgl.batch(graphs) |
| empty_globals = torch.empty((count, 0), dtype=torch.float32) |
| from gnn4colliders.models.root_gnn import ( |
| EdgeNetwork, |
| load_legacy_edge_network_state_dict, |
| ) |
|
|
| model = EdgeNetwork(graphs[0], empty_globals[:1], 64, 12, 4, 4, dropout=0.0) |
| payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) |
| load_legacy_edge_network_state_dict(model, payload) |
| model.train() |
| labels = torch.from_numpy(artifact.labels[:count]).to(torch.long) |
| weights = torch.from_numpy(artifact.weights[:count]).to(torch.float32) |
| optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) |
| histories = [] |
| for _ in range(args.steps): |
| optimizer.zero_grad() |
| logits = model(graph, None) |
| loss = _weighted_loss(logits, labels, weights) |
| loss.backward() |
| gradients = _flat_state(model, gradients=True) |
| optimizer.step() |
| histories.append(float(loss.detach())) |
| exp_avg, optimizer_steps = _flat_optimizer_state(model, optimizer, "exp_avg") |
| exp_avg_sq, _ = _flat_optimizer_state(model, optimizer, "exp_avg_sq") |
| np.savez_compressed( |
| args.output, |
| logits=logits.detach().cpu().numpy(), |
| loss=np.asarray(float(loss.detach())), |
| history=np.asarray(histories, dtype=np.float64), |
| gradients=gradients, |
| parameters=_flat_state(model), |
| optimizer_steps=optimizer_steps, |
| exp_avg=exp_avg, |
| exp_avg_sq=exp_avg_sq, |
| ) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
