File size: 3,535 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | #!/usr/bin/env python3
"""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())
|