TorchGraphMAE

GraphMAE2, rebuilt on PyTorch Geometric — so it actually installs and runs.

GraphMAE2 (WWW'23) is a strong masked autoencoder for graphs. The catch: the official code is built on DGL, which hasn't shipped a release in over two years. If you've tried to stand it up on a recent PyTorch + CUDA stack, you've probably already lost an afternoon to wheels that don't exist for your setup.

TorchGraphMAE is the same model, the same method, the same configs — re-implemented on PyTorch Geometric. There's no DGL anywhere: every graph op runs on edge_index with PyG's MessagePassing / scatter / softmax. If you live in PyG, this should feel like home.

Coming from the original repo? The DGL README is kept as README_original.md, and every file here maps 1:1 to the DGL file it replaced — see What's inside.

The smallest thing that works — pretrain and linear-probe on Cora. It auto-downloads the data and finishes in a couple of minutes on CPU or any GPU:

python main_full_batch_pyg.py --dataset cora --seeds 0 --use_cfg

That's the whole loop: mask the node features, train the autoencoder, then score the learned embeddings with a linear classifier. If this runs cleanly, your install is good and you can move on to the arxiv (minibatch) path.

GraphMAE2 is a masked feature autoencoder — it learns by hiding some node features and reconstructing them. That single fact decides what your graph needs to look like.

You need node features. This isn't optional — the features are the training signal, so they're what the model reconstructs. A structure-only graph (edges but no node vectors) has nothing to learn from. If that's all you have, synthesize features first (one-hot, degree stats, node2vec, …) and feed those in.

Everything else about the edges is read as a plain connection — just who links to whom:

One nuance on direction: the GAT layer itself does directed message passing (src → dst, softmax over each node's in-edges), so it can technically run on a directed edge_index. But every built-in loader symmetrizes the graph first (to_undirected / to_bidirected), so the default behavior is undirected. If you genuinely want directed propagation, drop the symmetrization in the loader and pass your own edge_index — the layer won't complain.

The short version: this is for homogeneous, attributed graphs — one node type, one edge type, real node-feature vectors. That's exactly the Planetoid / OGB node-classification setting GraphMAE2 was built for, and what this port reproduces.

Here's the honest version.

The model is a faithful, line-by-line port of the original — but I'm in the middle of re-running the large-scale benchmark, so I'd rather show you a blank than a number I can't stand behind yet.

For reference, here's what the original GraphMAE2 reports on the two datasets this port targets:

And here's where this port stands:

Why blank? I recently fixed a fidelity bug in ego-graph generation — the PPR diffusion wasn't adding self-loops the way the original does (details below). Any number I measured before that fix isn't an apples-to-apples comparison, so I pulled it instead of quietly leaving a stale one up. Fresh numbers land here once the clean run is done.

What I can point to today is that the compute is verified, not hoped-for:

• Cora inputs are bit-identical to DGL — same edges, labels, and splits, and after matching DGL's feature normalization, max|Δ| = 0 on the feature matrix itself.
• The GAT layer is a numerical equivalence, not an approximation (the exact DGL↔PyG mapping is spelled out below).
• Minibatch ego-graphs match dgl.batch exactly — edge-set difference of zero.
• Both entrypoints train and evaluate end-to-end without error.

So the foundation is solid; the headline accuracy number is just a re-run away.

Python 3.10, CUDA 12.4 wheels below (swap in your own CUDA build if different):

# PyTorch
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu124

# PyG core (pure Python) + optional accelerators (PyG falls back to native ops without them)
pip install torch-geometric
pip install pyg-lib torch-scatter torch-sparse torch-cluster \
    -f https://data.pyg.org/whl/torch-2.4.0+cu124.html

# the rest
pip install ogb scikit-learn pyyaml tqdm tensorboardX scipy

# ONLY needed for the ogbn-arxiv ego-graph step (skip if you're just running Cora):
SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True pip install localgraphclustering

No DGL. You will never be asked to install it.

There are two paths, depending on dataset size. Cora is the warm-up; arxiv is the real thing.

python main_full_batch_pyg.py --dataset cora --seeds 0 --use_cfg

Pretrains the masked autoencoder, then linear-probes. Data downloads itself via PyG Planetoid into $PYG_DATA_ROOT (default dataset/planetoid). Swap cora for citeseer or pubmed.

Bigger graph, so there's a one-time preprocessing step first.

Step 1 — build the LC ego-graphs (do this once; ~8 min for arxiv, CPU-only):

python -m datasets.localclustering_pyg --dataset ogbn-arxiv \
    --ego_size 256 --num_workers 4 --save_dir lc_ego_graphs

Step 2 — pretrain (needs a GPU; peaks around 30 GB, so reach for a ≥32 GB card):

python main_large_pyg.py --dataset ogbn-arxiv --use_cfg \
    --encoder gat --decoder gat --num_out_heads 1 --num_dec_layers 1 --residual \
    --sampling_method lc --seeds 0 \
    --ego_graph_file_path lc_ego_graphs/ogbn-arxiv-lc-ego-graphs-256.pt --data_dir dataset

Running on a cluster and want to split pretraining from evaluation? Set GM2_SKIP_EVAL=1 to stop right after pretraining (the checkpoint is saved), then evaluate later with --load_model --checkpoint_path <ckpt>. The arxiv data auto-downloads via OGB into $OGB_DATA_ROOT (default dataset).

Every PyG module is a drop-in replacement for one original DGL module (the DGL ones have been removed). If you know the original layout, this table is your map:

If you're auditing the port (or just curious how DGL graph ops become PyG ops), here's the substance.

The GAT layer is the only real graph computation, and it maps exactly. DGL splits the attention score as aᵀ[Wh_i‖Wh_j] = a_l·Wh_i + a_r·Wh_j, and each piece has a direct PyG counterpart:

This is a numerical equivalence, not an approximation — including the original's quirky residual rule (the residual is skipped when in_dim == out_dim * heads), which is preserved on purpose.

Cora features needed one careful fix. DGL's CoraGraphDataset hands back row-normalized features; PyG's Planetoid hands back raw binary ones, and GraphMAE2 doesn't rescale small datasets. So the loader applies NormalizeFeatures() (L1) — which makes the feature matrix bit-identical to DGL (max|Δ| = 0; edges/labels/splits already matched).

Minibatching rebuilds dgl.batch from scratch. dgl.batch([graph.subgraph(nodes)…]) becomes scipy-CSR induced-subgraph slicing (A[nodes][:, nodes], keeping the nodes[i]→i relabeling so each ego root stays at local index 0) plus node-offset edge_index concatenation. Verified to produce graphs identical to dgl.batch (edge-set difference = 0).

Ego-graph PPR runs with self-loops — and getting this wrong is exactly the fidelity bug mentioned up top. The original preprocess diffuses on to_bidirected().remove_self_loop().add_self_loop(), i.e. one self-loop per node. Those self-loops bump every node's degree by 1 (a lazy random walk) and feed the ppr × (1/deg) ranking that picks the top-ego_size neighbors, so they have to be there for the ego-graphs to match. localclustering_pyg.load_dataset_pyg now reproduces it (to_undirected → remove_self_loops → add_self_loops).

One free speedup. The original computes conductances via my_sweep_cut (O(ego_size²) sparse lookups) and then... never uses the result. Dropping it takes ego-graph generation from hours down to ~8 min on arxiv, with byte-for-byte identical output.

Two pieces of the original were left out on purpose, because nothing actually used them:

• models/gcn.py (the GCN backbone) — no config selects it; GAT is the default everywhere.
• datasets/saint_sampler.py — dead code in the original too; only lc sampling was ever wired up.

All credit for the method and the original implementation goes to GraphMAE2 by Hou et al. (THUDM): https://github.com/THUDM/GraphMAE2. This repo is just a framework port.

@inproceedings{hou2023graphmae2,
  title={GraphMAE2: A Decoding-Enhanced Masked Self-Supervised Graph Learner},
  author={Hou, Zhenyu and He, Yufei and Cen, Yukuo and Liu, Xiao and Dong, Yuxiao and Kharlamov, Evgeny and Tang, Jie},
  booktitle={WWW},
  year={2023}
}