File size: 2,081 Bytes
916755e de46a3c 916755e de46a3c 916755e de46a3c 916755e de46a3c 916755e de46a3c 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 | #!/usr/bin/env python3
"""Dump deterministic batch membership and aggregate sizes."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import torch
from validation.artifacts import load_artifact
from validation.forward import _graph
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--artifact", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--batch-size", type=int, default=8)
args = parser.parse_args()
artifact = load_artifact(args.artifact)
from gnn4colliders.data import (
EventMetadata,
GraphDataLoader,
GraphDataset,
GraphSample,
)
samples = []
for index in range(artifact.event_count):
samples.append(
GraphSample(
_graph(artifact, index),
torch.tensor(artifact.labels[index]),
None,
EventMetadata(
int(artifact.folds[index]),
float(artifact.weights[index]),
str(artifact.sample_id[index]),
{"index": index},
),
)
)
loader = GraphDataLoader(GraphDataset(samples), args.batch_size, shuffle=False)
batches = []
for batch in loader:
graph = batch.graph
labels = batch.labels
ids = list(batch.metadata.sample_id)
weights = batch.metadata.weight.tolist()
batches.append(
{
"sample_id": ids,
"num_graphs": len(ids),
"nodes": int(graph.num_nodes()),
"edges": int(graph.num_edges()),
"labels": labels.reshape(-1).tolist(),
"weights": weights,
}
)
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(
json.dumps({"batch_size": args.batch_size, "batches": batches}, indent=2) + "\n"
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|