Metadata-Version: 2.4
Name: graphssl
Version: 0.2.0
Summary: A modular Python library for Self-Supervised Learning on graphs, built on PyTorch and PyTorch Geometric.
Author-email: Simone Copetti <copetti.simone7@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/Simocop7/graphssl
Project-URL: Repository, https://github.com/Simocop7/graphssl
Project-URL: Bug Tracker, https://github.com/Simocop7/graphssl/issues
Project-URL: Documentation, https://simocop7.github.io/graphssl/
Project-URL: Changelog, https://github.com/Simocop7/graphssl/blob/main/CHANGELOG.md
Keywords: graph-learning,self-supervised-learning,gnn,pytorch-geometric,contrastive-learning,graph-neural-networks
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: torch_geometric>=2.3.0
Provides-Extra: viz
Requires-Dist: umap-learn; extra == "viz"
Requires-Dist: matplotlib; extra == "viz"
Requires-Dist: seaborn; extra == "viz"
Provides-Extra: full
Requires-Dist: faiss-cpu; extra == "full"
Requires-Dist: umap-learn; extra == "full"
Requires-Dist: matplotlib; extra == "full"
Requires-Dist: seaborn; extra == "full"
Requires-Dist: ogb; extra == "full"
Requires-Dist: pyyaml; extra == "full"
Provides-Extra: benchmark
Requires-Dist: ogb; extra == "benchmark"
Requires-Dist: pyyaml; extra == "benchmark"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Requires-Dist: pyyaml; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: mypy; extra == "dev"
Requires-Dist: pre-commit; extra == "dev"
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.5; extra == "docs"
Requires-Dist: mkdocstrings[python]>=0.25; extra == "docs"
Dynamic: license-file

# GraphSSL

[![Tests](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml/badge.svg)](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml)
[![Docs](https://github.com/Simocop7/graphssl/actions/workflows/docs.yml/badge.svg)](https://simocop7.github.io/graphssl/)
[![codecov](https://codecov.io/gh/Simocop7/graphssl/branch/main/graph/badge.svg)](https://codecov.io/gh/Simocop7/graphssl)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
[![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](pyproject.toml)
[![License: MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)

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

| 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.

---

## Installation

```bash
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:
> ```bash
> pip install torch --index-url https://download.pytorch.org/whl/cpu
> pip install graphssl
> ```

**Development:**
```bash
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](CONTRIBUTING.md) for the full dev workflow (lint, type checking, adding a model).

**Full documentation:** [simocop7.github.io/graphssl](https://simocop7.github.io/graphssl/) — per-model pages, architecture, and an auto-generated API reference. Build it locally with:
```bash
pip install -e ".[docs]"
mkdocs serve   # http://127.0.0.1:8000
```

---

## Quick Start

Every model uses the same constructor API:

```python
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:

```python
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/benchmark_planetoid.py` | Cora / CiteSeer / PubMed | Node classification, multi-seed | BGRL-only multi-seed demo; the cross-method numbers below come from `benchmarks/run_benchmark.py` |
| `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.

### Evaluate, save, fine-tune

```python
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](https://simocop7.github.io/graphssl/evaluation/) 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](https://github.com/Simocop7/graphssl/blob/main/CHANGELOG.md) 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**

| Method | Cora | CiteSeer | PubMed |
|---|---|---|---|
| *Untrained encoder* | 73.17 ± 0.99 | 57.69 ± 1.47 | 74.07 ± 1.60 |
| DGI | 78.12 ± 1.22 | 65.07 ± 1.44 | 78.60 ± 1.35 |
| GraphCL | **82.46 ± 1.08** | **69.59 ± 1.00** | 81.28 ± 0.77 |
| BGRL | 80.82 ± 1.23 | 67.05 ± 0.89 | 79.85 ± 1.10 |
| AFGRL | 79.18 ± 1.26 | 62.63 ± 2.01 | 78.43 ± 1.00 |
| VICReg | 81.06 ± 1.10 | 66.78 ± 1.45 | **81.59 ± 0.70** |
| Barlow Twins | 81.65 ± 0.93 | 68.95 ± 1.33 | 81.24 ± 0.72 |
| GraphDINO | 77.53 ± 1.22 | 60.16 ± 1.40 | 76.97 ± 2.08 |
| *Supervised* † | 72.67 ± 0.85 | 56.35 ± 1.52 | 74.12 ± 1.01 |
| *Supervised, std. recipe* † | 80.90 ± 0.76 | 67.72 ± 1.67 | 78.48 ± 0.92 |

**GIN encoder**

| Method | Cora | CiteSeer | PubMed |
|---|---|---|---|
| *Untrained encoder* | 41.70 ± 2.44 | 36.06 ± 1.57 | 56.60 ± 1.50 |
| DGI | 65.43 ± 3.03 | 49.18 ± 2.86 | 66.90 ± 3.87 |
| GraphCL | **81.31 ± 1.14** | 65.17 ± 1.44 | **80.36 ± 0.82** |
| BGRL | 63.77 ± 1.60 | 47.68 ± 2.30 | 68.52 ± 2.57 |
| AFGRL | 70.75 ± 1.75 | 49.32 ± 3.02 | 74.51 ± 1.11 |
| VICReg | 77.36 ± 1.88 | 63.52 ± 2.10 | 79.08 ± 1.27 |
| Barlow Twins | 79.33 ± 0.75 | **65.60 ± 1.62** | 77.67 ± 1.49 |
| GraphDINO | 62.53 ± 1.91 | 43.34 ± 3.27 | 70.38 ± 3.23 |
| *Supervised* † | 71.46 ± 1.83 | 51.21 ± 3.73 | 71.49 ± 1.70 |
| *Supervised, std. recipe* † | 75.83 ± 2.33 | 59.32 ± 4.51 | 71.96 ± 1.90 |

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.

```bash
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](https://simocop7.github.io/graphssl/benchmarks/) 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](https://simocop7.github.io/graphssl/benchmarks/#ogbn-arxiv-bgrl-with-the-protocol-of-its-paper).

```bash
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`:

```python
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:

```python
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:

```python
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:

```python
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](https://necst.it), Politecnico di Milano.
