GraphSSL

A modular Python library for Self-Supervised Learning on graphs, built on PyTorch and PyTorch Geometric, and the instrument of a controlled study of how graph SSL methods are evaluated. No Lightning, no Hydra — clean, readable training loops you can step through with a debugger.

Status: Alpha — all eight SSL objectives train end-to-end with every backbone and pass the test suite. The library is the instrument of an evaluation study of how much of a graph SSL result comes from the objective, and how much from the backbone, the probe and the training settings. What changed between releases is in the changelog.

---

Supported Methods

Method	Family	Paper
DGI	Mutual Information	Veličković et al., ICLR 2019
GraphCL	Contrastive (NT-Xent)	You et al., NeurIPS 2020
BGRL	Teacher-Student + EMA	Thakoor et al., ICLR 2022
AFGRL	Augmentation-Free Mining	Lee et al., AAAI 2022
VICReg	Variance-Invariance-Covariance	Bardes et al., ICLR 2022
Barlow Twins	Cross-Correlation	Zbontar et al., ICML 2021
GraphDINO	Self-Distillation	Adapted from Caron et al., ICCV 2021
GraphMAE	Masked Feature Reconstruction	Hou et al., KDD 2022

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

macOS and AFGRL: with the PyPI wheels, torch, faiss-cpu and scikit-learn each ship their own OpenMP library, and AFGRL ends with a segmentation fault. Make FAISS and scikit-learn use PyTorch's copy (reinstalling them undoes it):

SP=$(python -c "import site; print(site.getsitepackages()[0])")
for d in faiss/.dylibs sklearn/.dylibs; do
  mv "$SP/$d/libomp.dylib" "$SP/$d/libomp.dylib.bundled"
  ln -s "$SP/torch/lib/libomp.dylib" "$SP/$d/libomp.dylib"
done

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:

Script	Dataset	Task	Notes
examples/cora_bgrl.py	Cora	Node classification (7 classes)	Full-batch, CPU-friendly — fastest smoke test
examples/ogbn_arxiv_bgrl.py	ogbn-arxiv	Node classification (40 classes)	NeighborLoader mini-batch training
examples/zinc_bgrl.py	ZINC-12k	Molecular property regression	Categorical node/edge embeddings

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

This is the beginning of the README. Read the rest at the original. Read the full README ↗