| |
| """Bounded native WIRE-Performer run on the paper's MNIST graph benchmark.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import sys |
| import time |
| import types |
| from pathlib import Path |
|
|
| import networkx as nx |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch_geometric.datasets import GNNBenchmarkDataset |
| from torch_geometric.loader import DataLoader |
| from torch_geometric.nn import GCNConv, global_mean_pool |
|
|
|
|
| def import_pinned_graphrope(repo: Path): |
| graphgps = types.ModuleType("graphgps") |
| graphgps.__path__ = [str(repo / "graphgps")] |
| sys.modules["graphgps"] = graphgps |
| layer = types.ModuleType("graphgps.layer") |
| layer.__path__ = [str(repo / "graphgps" / "layer")] |
| sys.modules["graphgps.layer"] = layer |
| from graphgps.layer.graphrope import GraphRoPE |
|
|
| return GraphRoPE |
|
|
|
|
| def seed_all(seed: int): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
|
|
| def add_features(data, max_freqs: int): |
| n = int(data.num_nodes) |
| graph = nx.Graph() |
| graph.add_nodes_from(range(n)) |
| graph.add_edges_from(data.edge_index.t().tolist()) |
| adj = nx.to_numpy_array(graph, nodelist=range(n), dtype=float) |
| lap = np.diag(adj.sum(axis=1)) - adj |
| _, vecs = np.linalg.eigh(lap) |
| pe = vecs[:, 1 : 1 + max_freqs] |
| if pe.shape[1] < max_freqs: |
| pe = np.pad(pe, ((0, 0), (0, max_freqs - pe.shape[1]))) |
| pe = torch.tensor(pe, dtype=torch.float32) |
| data.x = torch.cat([data.x.float(), pe], dim=1) |
| data.t = pe |
| return data |
|
|
|
|
| class GPSMini(nn.Module): |
| def __init__(self, GraphRoPE, m: int, hidden: int = 32, layers: int = 3): |
| super().__init__() |
| self.m = m |
| self.input = nn.Linear(9, hidden) |
| self.local = nn.ModuleList([GCNConv(hidden, hidden) for _ in range(layers)]) |
| self.attn = nn.ModuleList([ |
| GraphRoPE( |
| k=max(1, m), d=hidden, num_heads=4, dropout=0.0, |
| enable=m > 0, init_omega="zero", attn_type="Linear", |
| ) |
| for _ in range(layers) |
| ]) |
| self.norm = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(layers)]) |
| self.ff = nn.ModuleList([ |
| nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden)) |
| for _ in range(layers) |
| ]) |
| self.head = nn.Linear(hidden, 10) |
|
|
| def forward(self, batch): |
| h = self.input(batch.x) |
| for local, attn, norm, ff in zip(self.local, self.attn, self.norm, self.ff): |
| h0 = h |
| h_local = local(h, batch.edge_index) |
| tmp = types.SimpleNamespace(x=h, batch=batch.batch) |
| if self.m > 0: |
| tmp.t = batch.t[:, : self.m] |
| h_attn = attn(tmp) |
| h = norm(h0 + h_local + h_attn) |
| h = h + ff(h) |
| return self.head(global_mean_pool(h, batch.batch)) |
|
|
|
|
| def run(GraphRoPE, root: Path, n_train: int, n_test: int, epochs: int, m: int, seed: int): |
| seed_all(seed) |
| train_ds = GNNBenchmarkDataset(root=str(root), name="MNIST", split="train") |
| test_ds = GNNBenchmarkDataset(root=str(root), name="MNIST", split="test") |
| train = [add_features(train_ds[i], 8) for i in range(n_train)] |
| test = [add_features(test_ds[i], 8) for i in range(n_test)] |
| train_loader = DataLoader(train, batch_size=16, shuffle=True) |
| test_loader = DataLoader(test, batch_size=32, shuffle=False) |
| model = GPSMini(GraphRoPE, m=m) |
| opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5) |
| best = 0.0 |
| for _ in range(epochs): |
| model.train() |
| for batch in train_loader: |
| opt.zero_grad(set_to_none=True) |
| loss = nn.functional.cross_entropy(model(batch), batch.y.view(-1)) |
| loss.backward() |
| opt.step() |
| model.eval() |
| good = total = 0 |
| with torch.no_grad(): |
| for batch in test_loader: |
| pred = model(batch).argmax(dim=-1) |
| good += int((pred == batch.y.view(-1)).sum()) |
| total += len(pred) |
| best = max(best, good / total) |
| return {"dataset": "MNIST", "attention": "Performer", "m": m, "seed": seed, |
| "train_graphs": n_train, "test_graphs": n_test, "epochs": epochs, |
| "best_test_accuracy": best, "parameters": sum(p.numel() for p in model.parameters())} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--repo", type=Path, required=True) |
| ap.add_argument("--data-root", type=Path, required=True) |
| ap.add_argument("--out", type=Path, required=True) |
| ap.add_argument("--train-graphs", type=int, default=256) |
| ap.add_argument("--test-graphs", type=int, default=256) |
| ap.add_argument("--epochs", type=int, default=3) |
| ap.add_argument("--seeds", type=int, nargs="+", default=[0, 1]) |
| args = ap.parse_args() |
| torch.set_num_threads(min(4, os.cpu_count() or 1)) |
| GraphRoPE = import_pinned_graphrope(args.repo) |
| start = time.time() |
| rows = [] |
| for m in [0, 8]: |
| for seed in args.seeds: |
| print(f"running MNIST Performer m={m} seed={seed}", flush=True) |
| rows.append(run(GraphRoPE, args.data_root, args.train_graphs, args.test_graphs, args.epochs, m, seed)) |
| print(f" best accuracy={rows[-1]['best_test_accuracy']:.6f}", flush=True) |
| out = { |
| "protocol": { |
| "paper": "arXiv:2509.22259v1, Section 4.3 / Table 3", |
| "official_repository": "https://github.com/cederikhoefs/Graph-RoPE", |
| "official_commit": "4ac067eb38272543b0cdd7591d630399ff37bce4", |
| "dataset": "PyG GNNBenchmarkDataset MNIST graph classification", |
| "architecture": "GCN local branch + official GraphRoPE Performer global branch, hidden=32, heads=4, layers=3", |
| "budget": {"train_graphs": args.train_graphs, "test_graphs": args.test_graphs, "epochs": args.epochs, "seeds": args.seeds}, |
| "baseline": "m=0 with the same Laplacian features as node inputs and WIRE disabled", |
| "wire": "m=8 spectral coordinates supplied to official GraphRoPE Performer", |
| }, |
| "rows": rows, |
| "runtime_seconds": time.time() - start, |
| } |
| args.out.parent.mkdir(parents=True, exist_ok=True) |
| args.out.write_text(json.dumps(out, indent=2) + "\n", encoding="utf-8") |
| print(f"wrote {args.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
