GraphSSL

A modular Python library for Self-Supervised Learning on graphs, built on PyTorch and PyTorch Geometric. No Lightning, no Hydra — clean, readable training loops you can step through with a debugger.

Status: Alpha — all models train end-to-end and pass tests; all 7 SSL methods benchmarked under a shared protocol on the citation networks (two backbones, see below) and on ZINC; BGRL reproduced on ogbn-arxiv with the protocol of its paper. OGB graph-level benchmarks are not wired up yet.

Supported Methods

Plus a Supervised baseline for comparison.

Encoder backbones: GCN, GIN, Graph Transformer — all swappable via one config line.

Installation

pip install graphssl

# optional extras
pip install "graphssl[viz]"        # UMAP + matplotlib
pip install "graphssl[benchmark]"  # ogb + pyyaml (for ZINC / ogbn-arxiv)
pip install "graphssl[full]"       # faiss-cpu + viz + benchmark

No GPU? pip install graphssl pulls in PyTorch's default CUDA-enabled build, which is large (1-2GB+ with bundled NVIDIA runtime libraries). On a CPU-only machine, install the CPU build first — it's a fraction of the size and installs much faster:

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install graphssl

Development:

git clone https://github.com/Simocop7/graphssl.git
cd graphssl
pip install -e ".[dev]"
pre-commit install   # ruff on every commit
pytest tests/ -v

All unit tests should pass. AFGRL tests are auto-skipped if faiss-cpu is not installed. See CONTRIBUTING.md for the full dev workflow (lint, type checking, adding a model).

Full documentation: simocop7.github.io/graphssl — per-model pages, architecture, and an auto-generated API reference. Build it locally with:

pip install -e ".[docs]"
mkdocs serve   # http://127.0.0.1:8000

Quick Start

Every model uses the same constructor API:

from graphssl.models import BGRL
from graphssl.training import DINOTrainer
from torch.optim import AdamW

config = {
    "encoder": {
        "name": "gcn",         # 'gcn', 'gin' or 'transformer'
        "hidden_dim": 256,
        "num_layers": 3,
        "norm_type": "batch",  # 'batch', 'layer', or 'none'
        "pool": False,         # False = node-level task
    },
    "augment": [
        {"name": "edge_drop", "p": 0.5},
        {"name": "feat_mask", "p": 0.1},
    ],
    "pred_hidden": 512,
    "ema_tau": 0.99,
    "ema_tau_end": 1.0,
    "total_steps": 0,
}

model = BGRL(config, in_channels=dataset.num_features)
optimizer = AdamW(model.student_parameters(), lr=1e-3)
trainer = DINOTrainer(device="cuda")
losses = trainer.train(model, loader, optimizer, num_epochs=1000)

Or load config from YAML:

from graphssl.config.load import load_config, build_model

cfg   = load_config("configs/ogbn_arxiv_bgrl.yaml")
model = build_model(cfg, in_channels=128)

See configs/ for reference YAML files and examples/ for full benchmark scripts:

Each example prints its result next to the same encoder left untrained: part of a frozen GNN's accuracy comes from the architecture alone.

Evaluate, save, fine-tune

from graphssl.config import load_model, save_model
from graphssl.core import pretrained_encoder
from graphssl.evaluation import LogRegEvaluator, effective_rank, extract_embeddings

z, y = extract_embeddings(model, datamodule)             # frozen embeddings
probe = LogRegEvaluator().evaluate(z, y, train_idx, val_idx, test_idx, num_classes=7)
probe["test_acc"], effective_rank(z)

save_model(model, "bgrl.pt", config, in_channels=dataset.num_features)
model = load_model("bgrl.pt")                            # rebuilt from the saved config
encoder = pretrained_encoder(model)                      # the encoder to fine-tune from

The evaluation guide has the full script: the probes (linear, kNN, ridge for regression), the untrained-encoder reference, and fine-tuning.

Upgrading from 0.1.0

0.2.0 changes behavior in several places; the changelog lists them all. The ones most likely to matter:

• Linear-probe numbers are not comparable with 0.1.0. LogRegEvaluator is now fitted to convergence on standardized features, with the L2 strength selected on validation; its arguments are keyword-only (lr and epochs are gone).
• A config field the encoder cannot use is an error, not silently dropped (edge_dim with gcn, for instance).
• Edge features reach the encoder. In 0.1.0 every model ignored edge_attr.
• GCN: num_layers is honoured (it was always 2), checkpoints have new state-dict keys, and eval-mode outputs differ after short trainings.
• GraphDINO has new defaults (ema_tau_base, teacher_temp), and DGI trains with one forward pass.

Benchmark results (citation networks)

All 7 SSL methods under one shared protocol, run with two backbones: a 2-layer GCN or GIN encoder (hidden_dim=256), full-batch training for 300 steps, AdamW (lr 5e-4), public Planetoid split. Linear-probe test accuracy (%) on the frozen embeddings (regularised logistic regression fitted to convergence, L2 strength selected on validation), mean ± std over 10 seeds; best SSL method per dataset in bold. Untrained encoder is the same encoder at initialization. The two rows marked † are supervised references trained on the labeled nodes (the shared protocol, and a standard recipe: dropout 0.5, Adam lr 0.01 + L2 5e-4, best-validation checkpoint and, with GCN, no normalization layer, as in Kipf & Welling's model), not ranked.

GCN encoder

GIN encoder

With GCN every objective improves on the untrained encoder and the top four are within 1.6–2.8 points of each other; with GIN the objectives separate, and BGRL, AFGRL, GraphDINO and DGI fall well behind.

D="--dataset Cora CiteSeer PubMed --seeds 10"
python benchmarks/run_benchmark.py $D --model all supervised supervised_reg                 # GIN
python benchmarks/run_benchmark.py $D --model all supervised supervised_reg --encoder gcn
python benchmarks/render_tables.py --summary   # rebuilds these tables from benchmarks/results/*.json

This is a deliberately untuned configuration — no per-method or per-dataset hyperparameter search — so it compares objectives at equal budget rather than reproducing each paper's best reported number. kNN results, the graph-level benchmark on ZINC (frozen probe and fine-tuning), per-seed metrics and full provenance (git commit, package versions) are in the benchmarks docs and benchmarks/results/. OGB graph-level benchmarks (ogbg-molhiv, ogbg-molpcba) are next.

Reproduction: BGRL on ogbn-arxiv

Run with the protocol of its own paper (Thakoor et al., ICLR 2022: full-graph training for 10,000 steps, 3-layer GCN with layer normalization and weight standardization) instead of the shared one, BGRL reaches 72.31 ± 0.14 validation / 71.32 ± 0.10 test accuracy over 3 seeds, against 72.53 ± 0.09 / 71.64 ± 0.12 in the paper (20 seeds). The same encoder left untrained gives 70.23 ± 0.18 / 69.53 ± 0.03 (paper: 69.90 ± 0.11 / 68.94 ± 0.15). Close to the paper, not an exact match. This is with the predictor of the public BGRL implementation (pred_norm: none, pred_activation: prelu); with the library's default predictor, which has a batch normalization, the same training ends 1.6 points below the untrained encoder (67.91 ± 0.78 test), so disable it when the encoder uses layer normalization. The protocol, the differences that remain and the predictor comparison are in the benchmarks docs.

python -u benchmarks/reproduce_bgrl_arxiv.py --seeds 3   # about 4 h per seed on an NVIDIA A2

Datasets with Categorical Features (ZINC)

ZINC has integer node and edge features. Use node_emb_num_classes and edge_emb_num_classes to replace linear projections with nn.Embedding:

config = {
    "encoder": {
        "name": "gin",
        "hidden_dim": 64,
        "num_layers": 4,
        "norm_type": "layer",
        "pool": True,
        "node_emb_num_classes": 28,  # 28 atom types
        "edge_dim": 64,
        "edge_emb_num_classes": 4,   # 4 bond types
    },
    ...
}
model = BGRL(config, in_channels=1)  # in_channels ignored when node_emb_num_classes is set

Composable Losses

Mix and weight multiple objectives without touching model code:

from graphssl.losses import CombinedLoss

loss_fn = CombinedLoss.from_config([
    {"name": "nt_xent",  "weight": 0.7, "tau": 0.5},
    {"name": "vicreg",   "weight": 0.3, "invariance": 25.0, "variance": 25.0, "covariance": 1.0},
])
loss = loss_fn(z1, z2)

Extending the Library

All major components are registered by name and can be replaced via config:

from graphssl.registry import ENCODERS, LOSSES, AUGMENTS

@ENCODERS.register("my_gat")
class GATEncoder(nn.Module):
    ...

# Then use it in any model config:
config = {"encoder": {"name": "my_gat", "hidden_dim": 128, ...}, ...}

Models expose compute_loss(batch), post_backward(), and post_step() hooks, so you can write your own training loop if needed:

for batch in loader:
    loss = model.compute_loss(batch)
    loss.backward()
    model.post_backward()
    optimizer.step()
    model.post_step()

Project Structure

src/graphssl/
├── core/          # BaseSSLModel, Callback, Registry, Protocol interfaces
├── config/        # Dataclass schemas + YAML loading (load_config, build_model)
├── registry/      # ENCODERS, HEADS, AUGMENTS, LOSSES, LOADERS, OBJECTIVES, DATASETS
├── encoders/      # GCN, GIN, Transformer (auto-registered)
├── models/        # All SSL models + Supervised baseline
├── losses/        # NT-Xent, DINO, VICReg, Barlow Twins, CosineRegression, CombinedLoss
├── augmentation/  # 7 transforms, compose(), MultiView
├── nn/            # MLP, Projector, Predictor, DINOHead, pooling
├── evaluation/    # LogRegEvaluator, KNNEvaluator, extract_embeddings
├── training/      # DINOTrainer + callbacks
├── data/          # DataModule (full-batch and NeighborLoader)
└── utils/         # update_ema_params, CosineDecayScheduler, CosineEMAScheduler, PositiveMiner

Key Design Choices

• Uniform API. Every model is ModelClass(config: Dict, in_channels: int). The config dict is validated by a dedicated dataclass in schema.py.
• Hook-driven training. The trainer calls post_backward() and post_step() at fixed points. Models implement these to freeze prototype gradients (GraphDINO) or update the EMA teacher (BGRL, AFGRL) — without subclassing the trainer.
• Protected-node augmentations. When using NeighborLoader for mini-batch training, seed nodes are shielded from node_drop so the loss remains valid.
• Registry pattern. All components are registered by name (@ENCODERS.register("gin")), so swapping architectures is a one-line config change.
• Composable losses. CombinedLoss enables weighted mixtures of any registered losses, configurable from YAML or programmatically.

Acknowledgments

Developed at NECSTLab, Politecnico di Milano.