| |
| """Compare legacy/new transfer-learning forward with a fixed binary head.""" |
|
|
| 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 |
|
|
| LEGACY_CONFIG = { |
| "module": "models.GCN", |
| "class": "Edge_Network", |
| "args": { |
| "hid_size": 64, |
| "out_size": 12, |
| "n_layers": 4, |
| "n_proc_steps": 4, |
| "dropout": 0.0, |
| }, |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument( |
| "--implementation", choices=("legacy", "rewrite"), required=True |
| ) |
| 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=96) |
| parser.add_argument("--trainable-backbone", action="store_true") |
| 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)] |
| batch_graph = dgl.batch(graphs) |
| empty_globals = torch.empty((1, 0), dtype=torch.float32) |
| if args.implementation == "legacy": |
| from models.GCN import Transferred_Learning_Finetuning |
|
|
| model = Transferred_Learning_Finetuning( |
| str(args.checkpoint), |
| LEGACY_CONFIG, |
| graphs[0], |
| empty_globals, |
| 64, |
| 1, |
| 4, |
| 4, |
| dropout=0.0, |
| frozen_pretraining=not args.trainable_backbone, |
| ) |
| classifier = model.classify |
| else: |
| from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork |
|
|
| backbone = EdgeNetwork(graphs[0], empty_globals, 64, 12, 4, 4, dropout=0.0) |
| payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) |
| from gnn4colliders.models.root_gnn import load_legacy_edge_network_state_dict |
|
|
| load_legacy_edge_network_state_dict(backbone, payload) |
| model = FineTunedEdgeNetwork( |
| backbone, 1, freeze_backbone=not args.trainable_backbone |
| ) |
| classifier = model.classifier |
| torch.manual_seed(20260818) |
| classifier.load_state_dict(torch.nn.Linear(64, 1).state_dict()) |
| model.eval() |
| with torch.inference_mode(): |
| logits = model(batch_graph, None).cpu().numpy() |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(args.output, logits=logits) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
