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 are benchmarked under a shared protocol on the citation networks (two backbones) and on ZINC, and BGRL is reproduced on ogbn-arxiv with the protocol of its paper (see Benchmarks). What changed since 0.1.0 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

Plus a Supervised baseline for comparison.

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

Why GraphSSL

Comparing SSL paradigms on graphs today usually means stitching together several papers' official repositories — different frameworks, different encoders, different evaluation protocols — before you can trust that a difference in reported accuracy reflects the method and not the harness. GraphSSL puts bootstrapping (BGRL, AFGRL), redundancy-reduction (VICReg, Barlow Twins), self-distillation (GraphDINO), contrastive (GraphCL) and mutual-information (DGI) methods behind one identical construction and training API, so swapping one for another is a one-line config change — see Architecture for how that's enforced, not just documented.

Quick start

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)

Continue with Getting Started for installation and a full worked example, or jump straight to a model page if you already know which method you want.

Project

Developed at NECSTLab, Politecnico di Milano. Source on GitHub, package on PyPI.