File size: 2,713 Bytes
916755e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env python3
"""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())
|