Upload node2vec_model.py with huggingface_hub

node2vec_model.py

@@ -0,0 +1,117 @@
+"""Node2Vec model for the edge_ML_expected_ge5 graph.
+
+Default run is a smoke check (builds the model, runs a few optimizer steps).
+Pass --epochs N to train, which writes embeddings to --out.
+"""
+
+import argparse
+import time
+
+import torch
+from torch_geometric.data import Data
+from torch_geometric.data.data import DataEdgeAttr, DataTensorAttr
+from torch_geometric.data.storage import BaseStorage, EdgeStorage, GlobalStorage
+from torch_geometric.nn import Node2Vec
+
+import os
+
+DATA_DIR = os.environ.get("DATA_DIR", ".")
+
+GRAPH = f"{DATA_DIR}/edge_ML_expected_ge5_pyg.pt"
+
+
+def load_graph(path: str = GRAPH) -> Data:
+    torch.serialization.add_safe_globals(
+        [Data, DataEdgeAttr, DataTensorAttr, BaseStorage, EdgeStorage, GlobalStorage]
+    )
+    return torch.load(path, weights_only=True)
+
+
+def build_model(data: Data, args: argparse.Namespace, device: torch.device) -> Node2Vec:
+    return Node2Vec(
+        data.edge_index,
+        embedding_dim=args.embedding_dim,
+        walk_length=args.walk_length,
+        context_size=args.context_size,
+        walks_per_node=args.walks_per_node,
+        num_negative_samples=args.num_negative_samples,
+        p=args.p,
+        q=args.q,
+        num_nodes=data.num_nodes,
+        sparse=True,  # pairs with SparseAdam; the embedding table is the only param
+    ).to(device)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--embedding-dim", type=int, default=128)
+    ap.add_argument("--walk-length", type=int, default=20)
+    ap.add_argument("--context-size", type=int, default=10)
+    ap.add_argument("--walks-per-node", type=int, default=10)
+    ap.add_argument("--num-negative-samples", type=int, default=1)
+    ap.add_argument("--p", type=float, default=1.0, help="return parameter")
+    ap.add_argument("--q", type=float, default=1.0, help="in-out parameter")
+    ap.add_argument("--batch-size", type=int, default=128)
+    ap.add_argument("--lr", type=float, default=0.01)
+    ap.add_argument("--num-workers", type=int, default=4)
+    ap.add_argument("--epochs", type=int, default=0, help="0 = smoke check only")
+    ap.add_argument("--steps", type=int, default=5, help="steps for the smoke check")
+    ap.add_argument("--out", default=f"{DATA_DIR}/edge_ML_expected_ge5_n2v.pt")
+    args = ap.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    data = load_graph()
+    model = build_model(data, args, device)
+
+    print(f"graph      : {data.num_nodes:,} nodes, {data.edge_index.size(1) // 2:,} undirected edges")
+    print(f"device     : {device}")
+    print(f"model      : {model}")
+    print(f"parameters : {sum(p.numel() for p in model.parameters()):,} "
+          f"({data.num_nodes:,} x {args.embedding_dim})")
+    print(f"walks      : length={args.walk_length} context={args.context_size} "
+          f"per_node={args.walks_per_node} p={args.p} q={args.q}")
+
+    loader = model.loader(batch_size=args.batch_size, shuffle=True,
+                          num_workers=args.num_workers)
+    optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=args.lr)
+    print(f"loader     : {len(loader):,} batches/epoch of {args.batch_size} seed nodes")
+
+    def run_epoch(max_steps: int | None = None) -> float:
+        model.train()
+        total, n = 0.0, 0
+        for i, (pos_rw, neg_rw) in enumerate(loader):
+            optimizer.zero_grad()
+            loss = model.loss(pos_rw.to(device), neg_rw.to(device))
+            loss.backward()
+            optimizer.step()
+            total, n = total + loss.item(), n + 1
+            if max_steps is not None and i + 1 >= max_steps:
+                break
+        return total / max(n, 1)
+
+    if args.epochs == 0:
+        t0 = time.perf_counter()
+        loss = run_epoch(max_steps=args.steps)
+        print(f"\nsmoke check: {args.steps} steps, mean loss {loss:.4f}, "
+              f"{time.perf_counter() - t0:.1f}s")
+        z = model()
+        print(f"embeddings : {tuple(z.shape)} {z.dtype} on {z.device}")
+        print("model built and training step verified; pass --epochs N to train")
+        return
+
+    for epoch in range(1, args.epochs + 1):
+        t0 = time.perf_counter()
+        loss = run_epoch()
+        print(f"epoch {epoch:>3}/{args.epochs}  loss {loss:.4f}  "
+              f"{time.perf_counter() - t0:.1f}s")
+
+    model.eval()
+    with torch.no_grad():
+        z = model().cpu()
+    torch.save({"embedding": z, "node_id": data.node_id,
+                "args": vars(args)}, args.out)
+    print(f"saved embeddings {tuple(z.shape)} -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()