| |
| """Dump deterministic batch membership and aggregate sizes for one implementation.""" |
|
|
| 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 |
|
|
|
|
| class _LegacyDataset(torch.utils.data.Dataset): |
| def __init__(self, artifact): |
| self.artifact = artifact |
|
|
| def __len__(self): |
| return self.artifact.event_count |
|
|
| def __getitem__(self, index): |
| graph = _graph(self.artifact, index) |
| tracking = torch.tensor( |
| [ |
| self.artifact.folds[index], |
| self.artifact.weights[index], |
| index, |
| ], |
| dtype=torch.float32, |
| ) |
| return ( |
| graph, |
| torch.tensor(self.artifact.labels[index]), |
| tracking, |
| torch.empty(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("--output", type=Path, required=True) |
| parser.add_argument("--batch-size", type=int, default=8) |
| args = parser.parse_args() |
| artifact = load_artifact(args.artifact) |
| if args.implementation == "legacy": |
| from dgl.dataloading import GraphDataLoader |
|
|
| loader = GraphDataLoader( |
| _LegacyDataset(artifact), |
| batch_size=args.batch_size, |
| shuffle=False, |
| drop_last=False, |
| num_workers=0, |
| ) |
| else: |
| 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: |
| if args.implementation == "legacy": |
| graph, labels, tracking, _ = batch |
| indices = tracking[:, 2].to(torch.long).tolist() |
| ids = [str(artifact.sample_id[index]) for index in indices] |
| weights = tracking[:, 1].tolist() |
| else: |
| 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()) |
|
|
