graphssl 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (74) hide show
  1. graphssl-0.1.0/LICENSE +21 -0
  2. graphssl-0.1.0/PKG-INFO +273 -0
  3. graphssl-0.1.0/README.md +224 -0
  4. graphssl-0.1.0/pyproject.toml +118 -0
  5. graphssl-0.1.0/setup.cfg +4 -0
  6. graphssl-0.1.0/src/graphssl/__init__.py +63 -0
  7. graphssl-0.1.0/src/graphssl/augmentation/__init__.py +11 -0
  8. graphssl-0.1.0/src/graphssl/augmentation/compose.py +60 -0
  9. graphssl-0.1.0/src/graphssl/augmentation/functional.py +127 -0
  10. graphssl-0.1.0/src/graphssl/augmentation/transforms.py +84 -0
  11. graphssl-0.1.0/src/graphssl/config/__init__.py +14 -0
  12. graphssl-0.1.0/src/graphssl/config/load.py +114 -0
  13. graphssl-0.1.0/src/graphssl/config/schema.py +413 -0
  14. graphssl-0.1.0/src/graphssl/core/__init__.py +14 -0
  15. graphssl-0.1.0/src/graphssl/core/augmentation.py +16 -0
  16. graphssl-0.1.0/src/graphssl/core/callback.py +44 -0
  17. graphssl-0.1.0/src/graphssl/core/encoder.py +18 -0
  18. graphssl-0.1.0/src/graphssl/core/model.py +58 -0
  19. graphssl-0.1.0/src/graphssl/core/registry.py +47 -0
  20. graphssl-0.1.0/src/graphssl/data/__init__.py +1 -0
  21. graphssl-0.1.0/src/graphssl/data/datamodule.py +100 -0
  22. graphssl-0.1.0/src/graphssl/encoders/__init__.py +3 -0
  23. graphssl-0.1.0/src/graphssl/encoders/gcn.py +104 -0
  24. graphssl-0.1.0/src/graphssl/encoders/gin.py +181 -0
  25. graphssl-0.1.0/src/graphssl/encoders/transformer.py +206 -0
  26. graphssl-0.1.0/src/graphssl/evaluation/__init__.py +3 -0
  27. graphssl-0.1.0/src/graphssl/evaluation/knn.py +52 -0
  28. graphssl-0.1.0/src/graphssl/evaluation/linear_probe.py +116 -0
  29. graphssl-0.1.0/src/graphssl/evaluation/visualization.py +206 -0
  30. graphssl-0.1.0/src/graphssl/losses/__init__.py +6 -0
  31. graphssl-0.1.0/src/graphssl/losses/barlow.py +35 -0
  32. graphssl-0.1.0/src/graphssl/losses/combined.py +67 -0
  33. graphssl-0.1.0/src/graphssl/losses/dino.py +38 -0
  34. graphssl-0.1.0/src/graphssl/losses/nt_xent.py +40 -0
  35. graphssl-0.1.0/src/graphssl/losses/regression.py +44 -0
  36. graphssl-0.1.0/src/graphssl/losses/vicreg.py +56 -0
  37. graphssl-0.1.0/src/graphssl/models/__init__.py +14 -0
  38. graphssl-0.1.0/src/graphssl/models/afgrl.py +133 -0
  39. graphssl-0.1.0/src/graphssl/models/barlow_twins.py +68 -0
  40. graphssl-0.1.0/src/graphssl/models/bgrl.py +110 -0
  41. graphssl-0.1.0/src/graphssl/models/dgi.py +85 -0
  42. graphssl-0.1.0/src/graphssl/models/graphcl.py +68 -0
  43. graphssl-0.1.0/src/graphssl/models/graphdino.py +182 -0
  44. graphssl-0.1.0/src/graphssl/models/supervised.py +48 -0
  45. graphssl-0.1.0/src/graphssl/models/vicreg.py +72 -0
  46. graphssl-0.1.0/src/graphssl/nn/__init__.py +4 -0
  47. graphssl-0.1.0/src/graphssl/nn/dino_head.py +99 -0
  48. graphssl-0.1.0/src/graphssl/nn/mlp.py +52 -0
  49. graphssl-0.1.0/src/graphssl/nn/norm.py +24 -0
  50. graphssl-0.1.0/src/graphssl/nn/pooling.py +18 -0
  51. graphssl-0.1.0/src/graphssl/registry/__init__.py +1 -0
  52. graphssl-0.1.0/src/graphssl/registry/registry.py +9 -0
  53. graphssl-0.1.0/src/graphssl/training/__init__.py +2 -0
  54. graphssl-0.1.0/src/graphssl/training/callbacks.py +138 -0
  55. graphssl-0.1.0/src/graphssl/training/trainer.py +108 -0
  56. graphssl-0.1.0/src/graphssl/utils/__init__.py +3 -0
  57. graphssl-0.1.0/src/graphssl/utils/ema.py +14 -0
  58. graphssl-0.1.0/src/graphssl/utils/positive_miner.py +119 -0
  59. graphssl-0.1.0/src/graphssl/utils/schedulers.py +55 -0
  60. graphssl-0.1.0/src/graphssl.egg-info/PKG-INFO +273 -0
  61. graphssl-0.1.0/src/graphssl.egg-info/SOURCES.txt +72 -0
  62. graphssl-0.1.0/src/graphssl.egg-info/dependency_links.txt +1 -0
  63. graphssl-0.1.0/src/graphssl.egg-info/requires.txt +29 -0
  64. graphssl-0.1.0/src/graphssl.egg-info/top_level.txt +1 -0
  65. graphssl-0.1.0/tests/test_afgrl.py +129 -0
  66. graphssl-0.1.0/tests/test_barlow_twins.py +108 -0
  67. graphssl-0.1.0/tests/test_bgrl.py +141 -0
  68. graphssl-0.1.0/tests/test_dgi.py +105 -0
  69. graphssl-0.1.0/tests/test_evaluation.py +69 -0
  70. graphssl-0.1.0/tests/test_graphcl.py +97 -0
  71. graphssl-0.1.0/tests/test_graphdino.py +257 -0
  72. graphssl-0.1.0/tests/test_new_features.py +255 -0
  73. graphssl-0.1.0/tests/test_supervised.py +107 -0
  74. graphssl-0.1.0/tests/test_vicreg.py +106 -0
graphssl-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Simone Copetti
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,273 @@
1
+ Metadata-Version: 2.4
2
+ Name: graphssl
3
+ Version: 0.1.0
4
+ Summary: A modular Python library for Self-Supervised Learning on graphs, built on PyTorch and PyTorch Geometric.
5
+ Author-email: Simone Copetti <copetti.simone7@gmail.com>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/Simocop7/graphssl
8
+ Project-URL: Repository, https://github.com/Simocop7/graphssl
9
+ Project-URL: Bug Tracker, https://github.com/Simocop7/graphssl/issues
10
+ Project-URL: Documentation, https://github.com/Simocop7/graphssl#readme
11
+ Keywords: graph-learning,self-supervised-learning,gnn,pytorch-geometric,contrastive-learning,graph-neural-networks
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
20
+ Requires-Python: >=3.10
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: torch>=2.0.0
24
+ Requires-Dist: torch_geometric>=2.3.0
25
+ Provides-Extra: viz
26
+ Requires-Dist: umap-learn; extra == "viz"
27
+ Requires-Dist: matplotlib; extra == "viz"
28
+ Requires-Dist: seaborn; extra == "viz"
29
+ Provides-Extra: full
30
+ Requires-Dist: faiss-cpu; extra == "full"
31
+ Requires-Dist: umap-learn; extra == "full"
32
+ Requires-Dist: matplotlib; extra == "full"
33
+ Requires-Dist: seaborn; extra == "full"
34
+ Requires-Dist: ogb; extra == "full"
35
+ Requires-Dist: pyyaml; extra == "full"
36
+ Provides-Extra: benchmark
37
+ Requires-Dist: ogb; extra == "benchmark"
38
+ Requires-Dist: pyyaml; extra == "benchmark"
39
+ Provides-Extra: dev
40
+ Requires-Dist: pytest; extra == "dev"
41
+ Requires-Dist: pytest-cov; extra == "dev"
42
+ Requires-Dist: build; extra == "dev"
43
+ Requires-Dist: twine; extra == "dev"
44
+ Requires-Dist: pyyaml; extra == "dev"
45
+ Requires-Dist: ruff; extra == "dev"
46
+ Requires-Dist: mypy; extra == "dev"
47
+ Requires-Dist: pre-commit; extra == "dev"
48
+ Dynamic: license-file
49
+
50
+ # GraphSSL
51
+
52
+ [![Tests](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml/badge.svg)](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml)
53
+ [![codecov](https://codecov.io/gh/Simocop7/graphssl/branch/main/graph/badge.svg)](https://codecov.io/gh/Simocop7/graphssl)
54
+ [![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)
55
+ [![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](pyproject.toml)
56
+ [![License: MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)
57
+
58
+ 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.
59
+
60
+ > **Status:** Alpha — all models train end-to-end and pass tests; citation-network benchmarks validated (see below), large-scale (OGB) benchmarks in progress.
61
+
62
+ ---
63
+
64
+ ## Supported Methods
65
+
66
+ | Method | Family | Paper |
67
+ |---|---|---|
68
+ | **DGI** | Mutual Information | Veličković et al., ICLR 2019 |
69
+ | **GraphCL** | Contrastive (NT-Xent) | You et al., NeurIPS 2020 |
70
+ | **BGRL** | Teacher-Student + EMA | Thakoor et al., ICLR 2022 |
71
+ | **AFGRL** | Augmentation-Free Mining | Lee et al., AAAI 2022 |
72
+ | **VICReg** | Variance-Invariance-Covariance | Bardes et al., ICLR 2022 |
73
+ | **Barlow Twins** | Cross-Correlation | Zbontar et al., ICML 2021 |
74
+ | **GraphDINO** | Self-Distillation | Adapted from Caron et al., ICCV 2021 |
75
+
76
+ Plus a **Supervised** baseline for comparison.
77
+
78
+ Encoder backbones: **GCN**, **GIN**, **Graph Transformer** — all swappable via one config line.
79
+
80
+ ---
81
+
82
+ ## Installation
83
+
84
+ ```bash
85
+ pip install graphssl
86
+
87
+ # optional extras
88
+ pip install "graphssl[viz]" # UMAP + matplotlib
89
+ pip install "graphssl[benchmark]" # ogb + pyyaml (for ZINC / ogbn-arxiv)
90
+ pip install "graphssl[full]" # faiss-cpu + viz + benchmark
91
+ ```
92
+
93
+ **Development:**
94
+ ```bash
95
+ git clone https://github.com/Simocop7/graphssl.git
96
+ cd graphssl
97
+ pip install -e ".[dev]"
98
+ pre-commit install # ruff on every commit
99
+ pytest tests/ -v
100
+ ```
101
+
102
+ All 107 unit tests should pass. AFGRL tests are auto-skipped if `faiss-cpu` is not installed.
103
+ See [CONTRIBUTING.md](CONTRIBUTING.md) for the full dev workflow (lint, type checking, adding a model).
104
+
105
+ ---
106
+
107
+ ## Quick Start
108
+
109
+ Every model uses the same constructor API:
110
+
111
+ ```python
112
+ from graphssl.models import BGRL
113
+ from graphssl.training import DINOTrainer
114
+ from torch.optim import AdamW
115
+
116
+ config = {
117
+ "encoder": {
118
+ "name": "gin",
119
+ "hidden_dim": 256,
120
+ "num_layers": 3,
121
+ "norm_type": "batch", # 'batch', 'layer', or 'none'
122
+ "pool": False, # False = node-level task
123
+ },
124
+ "augment": [
125
+ {"name": "edge_drop", "p": 0.5},
126
+ {"name": "feat_mask", "p": 0.1},
127
+ ],
128
+ "pred_hidden": 512,
129
+ "ema_tau": 0.99,
130
+ "ema_tau_end": 1.0,
131
+ "total_steps": 0,
132
+ }
133
+
134
+ model = BGRL(config, in_channels=dataset.num_features)
135
+ optimizer = AdamW(model.student_parameters(), lr=1e-3)
136
+ trainer = DINOTrainer(device="cuda")
137
+ losses = trainer.train(model, loader, optimizer, num_epochs=1000)
138
+ ```
139
+
140
+ Or load config from YAML:
141
+
142
+ ```python
143
+ from graphssl.config.load import load_config, build_model
144
+
145
+ cfg = load_config("configs/ogbn_arxiv_bgrl.yaml")
146
+ model = build_model(cfg, in_channels=128)
147
+ ```
148
+
149
+ See `configs/` for reference YAML files and `examples/` for full benchmark scripts:
150
+
151
+ | Script | Dataset | Task | Notes |
152
+ |---|---|---|---|
153
+ | `examples/cora_bgrl.py` | Cora | Node classification (7 classes) | Full-batch, CPU-friendly — fastest smoke test |
154
+ | `examples/benchmark_planetoid.py` | Cora / CiteSeer / PubMed | Node classification, multi-seed | Reproduces the results below (`--seeds`, `--epochs`) |
155
+ | `examples/ogbn_arxiv_bgrl.py` | ogbn-arxiv | Node classification (40 classes) | `NeighborLoader` mini-batch training |
156
+ | `examples/zinc_bgrl.py` | ZINC-12k | Molecular property regression | Categorical node/edge embeddings |
157
+
158
+ ### Benchmark results (citation networks)
159
+
160
+ BGRL, GIN-2L encoder (`hidden_dim=256`), full-batch training for 300 steps, public Planetoid split, mean ± std over 10 seeds (`python examples/benchmark_planetoid.py --dataset <name> --seeds 10`):
161
+
162
+ | Dataset | Linear probe (test) | KNN k=5 (test) |
163
+ |---|---|---|
164
+ | Cora | 62.39 ± 3.71 | 58.61 ± 3.76 |
165
+ | CiteSeer | 50.08 ± 1.58 | 43.07 ± 3.61 |
166
+ | PubMed | 69.18 ± 2.32 | 66.59 ± 2.88 |
167
+
168
+ This is a deliberately lightweight, untuned configuration (no per-dataset hyperparameter search) meant to validate that the training/evaluation pipeline is correct end-to-end — not a state-of-the-art claim. See `paper.tex` for a methodological comparison against the original BGRL paper's own numbers on these datasets. ogbn-arxiv and OGB graph-level benchmarks (ogbg-molhiv, ogbg-molpcba) require larger compute and are planned on dedicated infrastructure.
169
+
170
+ ---
171
+
172
+ ## Datasets with Categorical Features (ZINC)
173
+
174
+ ZINC has integer node and edge features. Use `node_emb_num_classes` and `edge_emb_num_classes`
175
+ to replace linear projections with `nn.Embedding`:
176
+
177
+ ```python
178
+ config = {
179
+ "encoder": {
180
+ "name": "gin",
181
+ "hidden_dim": 64,
182
+ "num_layers": 4,
183
+ "norm_type": "layer",
184
+ "pool": True,
185
+ "node_emb_num_classes": 28, # 28 atom types
186
+ "edge_dim": 64,
187
+ "edge_emb_num_classes": 4, # 4 bond types
188
+ },
189
+ ...
190
+ }
191
+ model = BGRL(config, in_channels=1) # in_channels ignored when node_emb_num_classes is set
192
+ ```
193
+
194
+ ---
195
+
196
+ ## Composable Losses
197
+
198
+ Mix and weight multiple objectives without touching model code:
199
+
200
+ ```python
201
+ from graphssl.losses import CombinedLoss
202
+
203
+ loss_fn = CombinedLoss.from_config([
204
+ {"name": "nt_xent", "weight": 0.7, "tau": 0.5},
205
+ {"name": "vicreg", "weight": 0.3, "invariance": 25.0, "variance": 25.0, "covariance": 1.0},
206
+ ])
207
+ loss = loss_fn(z1, z2)
208
+ ```
209
+
210
+ ---
211
+
212
+ ## Extending the Library
213
+
214
+ All major components are registered by name and can be replaced via config:
215
+
216
+ ```python
217
+ from graphssl.registry import ENCODERS, LOSSES, AUGMENTS
218
+
219
+ @ENCODERS.register("my_gat")
220
+ class GATEncoder(nn.Module):
221
+ ...
222
+
223
+ # Then use it in any model config:
224
+ config = {"encoder": {"name": "my_gat", "hidden_dim": 128, ...}, ...}
225
+ ```
226
+
227
+ Models expose `compute_loss(batch)`, `post_backward()`, and `post_step()` hooks,
228
+ so you can write your own training loop if needed:
229
+
230
+ ```python
231
+ for batch in loader:
232
+ loss = model.compute_loss(batch)
233
+ loss.backward()
234
+ model.post_backward()
235
+ optimizer.step()
236
+ model.post_step()
237
+ ```
238
+
239
+ ---
240
+
241
+ ## Project Structure
242
+
243
+ ```
244
+ src/graphssl/
245
+ ├── core/ # BaseSSLModel, Callback, Registry, Protocol interfaces
246
+ ├── config/ # Dataclass schemas + YAML loading (load_config, build_model)
247
+ ├── registry/ # ENCODERS, HEADS, AUGMENTS, LOSSES, LOADERS, OBJECTIVES, DATASETS
248
+ ├── encoders/ # GCN, GIN, Transformer (auto-registered)
249
+ ├── models/ # All SSL models + Supervised baseline
250
+ ├── losses/ # NT-Xent, DINO, VICReg, Barlow Twins, CosineRegression, CombinedLoss
251
+ ├── augmentation/ # 7 transforms, compose(), MultiView
252
+ ├── nn/ # MLP, Projector, Predictor, DINOHead, pooling
253
+ ├── evaluation/ # LogRegEvaluator, KNNEvaluator, extract_embeddings
254
+ ├── training/ # DINOTrainer + callbacks
255
+ ├── data/ # DataModule (full-batch and NeighborLoader)
256
+ └── utils/ # update_ema_params, CosineDecayScheduler, CosineEMAScheduler, PositiveMiner
257
+ ```
258
+
259
+ ---
260
+
261
+ ## Key Design Choices
262
+
263
+ - **Uniform API.** Every model is `ModelClass(config: Dict, in_channels: int)`. The config dict is validated by a dedicated dataclass in `schema.py`.
264
+ - **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.
265
+ - **Protected-node augmentations.** When using `NeighborLoader` for mini-batch training, seed nodes are shielded from `node_drop` so the loss remains valid.
266
+ - **Registry pattern.** All components are registered by name (`@ENCODERS.register("gin")`), so swapping architectures is a one-line config change.
267
+ - **Composable losses.** `CombinedLoss` enables weighted mixtures of any registered losses, configurable from YAML or programmatically.
268
+
269
+ ---
270
+
271
+ ## Acknowledgments
272
+
273
+ Developed at [NECSTLab](https://necst.it), Politecnico di Milano.
@@ -0,0 +1,224 @@
1
+ # GraphSSL
2
+
3
+ [![Tests](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml/badge.svg)](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml)
4
+ [![codecov](https://codecov.io/gh/Simocop7/graphssl/branch/main/graph/badge.svg)](https://codecov.io/gh/Simocop7/graphssl)
5
+ [![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)
6
+ [![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](pyproject.toml)
7
+ [![License: MIT](https://img.shields.io/badge/license-MIT-green)](LICENSE)
8
+
9
+ 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.
10
+
11
+ > **Status:** Alpha — all models train end-to-end and pass tests; citation-network benchmarks validated (see below), large-scale (OGB) benchmarks in progress.
12
+
13
+ ---
14
+
15
+ ## Supported Methods
16
+
17
+ | Method | Family | Paper |
18
+ |---|---|---|
19
+ | **DGI** | Mutual Information | Veličković et al., ICLR 2019 |
20
+ | **GraphCL** | Contrastive (NT-Xent) | You et al., NeurIPS 2020 |
21
+ | **BGRL** | Teacher-Student + EMA | Thakoor et al., ICLR 2022 |
22
+ | **AFGRL** | Augmentation-Free Mining | Lee et al., AAAI 2022 |
23
+ | **VICReg** | Variance-Invariance-Covariance | Bardes et al., ICLR 2022 |
24
+ | **Barlow Twins** | Cross-Correlation | Zbontar et al., ICML 2021 |
25
+ | **GraphDINO** | Self-Distillation | Adapted from Caron et al., ICCV 2021 |
26
+
27
+ Plus a **Supervised** baseline for comparison.
28
+
29
+ Encoder backbones: **GCN**, **GIN**, **Graph Transformer** — all swappable via one config line.
30
+
31
+ ---
32
+
33
+ ## Installation
34
+
35
+ ```bash
36
+ pip install graphssl
37
+
38
+ # optional extras
39
+ pip install "graphssl[viz]" # UMAP + matplotlib
40
+ pip install "graphssl[benchmark]" # ogb + pyyaml (for ZINC / ogbn-arxiv)
41
+ pip install "graphssl[full]" # faiss-cpu + viz + benchmark
42
+ ```
43
+
44
+ **Development:**
45
+ ```bash
46
+ git clone https://github.com/Simocop7/graphssl.git
47
+ cd graphssl
48
+ pip install -e ".[dev]"
49
+ pre-commit install # ruff on every commit
50
+ pytest tests/ -v
51
+ ```
52
+
53
+ All 107 unit tests should pass. AFGRL tests are auto-skipped if `faiss-cpu` is not installed.
54
+ See [CONTRIBUTING.md](CONTRIBUTING.md) for the full dev workflow (lint, type checking, adding a model).
55
+
56
+ ---
57
+
58
+ ## Quick Start
59
+
60
+ Every model uses the same constructor API:
61
+
62
+ ```python
63
+ from graphssl.models import BGRL
64
+ from graphssl.training import DINOTrainer
65
+ from torch.optim import AdamW
66
+
67
+ config = {
68
+ "encoder": {
69
+ "name": "gin",
70
+ "hidden_dim": 256,
71
+ "num_layers": 3,
72
+ "norm_type": "batch", # 'batch', 'layer', or 'none'
73
+ "pool": False, # False = node-level task
74
+ },
75
+ "augment": [
76
+ {"name": "edge_drop", "p": 0.5},
77
+ {"name": "feat_mask", "p": 0.1},
78
+ ],
79
+ "pred_hidden": 512,
80
+ "ema_tau": 0.99,
81
+ "ema_tau_end": 1.0,
82
+ "total_steps": 0,
83
+ }
84
+
85
+ model = BGRL(config, in_channels=dataset.num_features)
86
+ optimizer = AdamW(model.student_parameters(), lr=1e-3)
87
+ trainer = DINOTrainer(device="cuda")
88
+ losses = trainer.train(model, loader, optimizer, num_epochs=1000)
89
+ ```
90
+
91
+ Or load config from YAML:
92
+
93
+ ```python
94
+ from graphssl.config.load import load_config, build_model
95
+
96
+ cfg = load_config("configs/ogbn_arxiv_bgrl.yaml")
97
+ model = build_model(cfg, in_channels=128)
98
+ ```
99
+
100
+ See `configs/` for reference YAML files and `examples/` for full benchmark scripts:
101
+
102
+ | Script | Dataset | Task | Notes |
103
+ |---|---|---|---|
104
+ | `examples/cora_bgrl.py` | Cora | Node classification (7 classes) | Full-batch, CPU-friendly — fastest smoke test |
105
+ | `examples/benchmark_planetoid.py` | Cora / CiteSeer / PubMed | Node classification, multi-seed | Reproduces the results below (`--seeds`, `--epochs`) |
106
+ | `examples/ogbn_arxiv_bgrl.py` | ogbn-arxiv | Node classification (40 classes) | `NeighborLoader` mini-batch training |
107
+ | `examples/zinc_bgrl.py` | ZINC-12k | Molecular property regression | Categorical node/edge embeddings |
108
+
109
+ ### Benchmark results (citation networks)
110
+
111
+ BGRL, GIN-2L encoder (`hidden_dim=256`), full-batch training for 300 steps, public Planetoid split, mean ± std over 10 seeds (`python examples/benchmark_planetoid.py --dataset <name> --seeds 10`):
112
+
113
+ | Dataset | Linear probe (test) | KNN k=5 (test) |
114
+ |---|---|---|
115
+ | Cora | 62.39 ± 3.71 | 58.61 ± 3.76 |
116
+ | CiteSeer | 50.08 ± 1.58 | 43.07 ± 3.61 |
117
+ | PubMed | 69.18 ± 2.32 | 66.59 ± 2.88 |
118
+
119
+ This is a deliberately lightweight, untuned configuration (no per-dataset hyperparameter search) meant to validate that the training/evaluation pipeline is correct end-to-end — not a state-of-the-art claim. See `paper.tex` for a methodological comparison against the original BGRL paper's own numbers on these datasets. ogbn-arxiv and OGB graph-level benchmarks (ogbg-molhiv, ogbg-molpcba) require larger compute and are planned on dedicated infrastructure.
120
+
121
+ ---
122
+
123
+ ## Datasets with Categorical Features (ZINC)
124
+
125
+ ZINC has integer node and edge features. Use `node_emb_num_classes` and `edge_emb_num_classes`
126
+ to replace linear projections with `nn.Embedding`:
127
+
128
+ ```python
129
+ config = {
130
+ "encoder": {
131
+ "name": "gin",
132
+ "hidden_dim": 64,
133
+ "num_layers": 4,
134
+ "norm_type": "layer",
135
+ "pool": True,
136
+ "node_emb_num_classes": 28, # 28 atom types
137
+ "edge_dim": 64,
138
+ "edge_emb_num_classes": 4, # 4 bond types
139
+ },
140
+ ...
141
+ }
142
+ model = BGRL(config, in_channels=1) # in_channels ignored when node_emb_num_classes is set
143
+ ```
144
+
145
+ ---
146
+
147
+ ## Composable Losses
148
+
149
+ Mix and weight multiple objectives without touching model code:
150
+
151
+ ```python
152
+ from graphssl.losses import CombinedLoss
153
+
154
+ loss_fn = CombinedLoss.from_config([
155
+ {"name": "nt_xent", "weight": 0.7, "tau": 0.5},
156
+ {"name": "vicreg", "weight": 0.3, "invariance": 25.0, "variance": 25.0, "covariance": 1.0},
157
+ ])
158
+ loss = loss_fn(z1, z2)
159
+ ```
160
+
161
+ ---
162
+
163
+ ## Extending the Library
164
+
165
+ All major components are registered by name and can be replaced via config:
166
+
167
+ ```python
168
+ from graphssl.registry import ENCODERS, LOSSES, AUGMENTS
169
+
170
+ @ENCODERS.register("my_gat")
171
+ class GATEncoder(nn.Module):
172
+ ...
173
+
174
+ # Then use it in any model config:
175
+ config = {"encoder": {"name": "my_gat", "hidden_dim": 128, ...}, ...}
176
+ ```
177
+
178
+ Models expose `compute_loss(batch)`, `post_backward()`, and `post_step()` hooks,
179
+ so you can write your own training loop if needed:
180
+
181
+ ```python
182
+ for batch in loader:
183
+ loss = model.compute_loss(batch)
184
+ loss.backward()
185
+ model.post_backward()
186
+ optimizer.step()
187
+ model.post_step()
188
+ ```
189
+
190
+ ---
191
+
192
+ ## Project Structure
193
+
194
+ ```
195
+ src/graphssl/
196
+ ├── core/ # BaseSSLModel, Callback, Registry, Protocol interfaces
197
+ ├── config/ # Dataclass schemas + YAML loading (load_config, build_model)
198
+ ├── registry/ # ENCODERS, HEADS, AUGMENTS, LOSSES, LOADERS, OBJECTIVES, DATASETS
199
+ ├── encoders/ # GCN, GIN, Transformer (auto-registered)
200
+ ├── models/ # All SSL models + Supervised baseline
201
+ ├── losses/ # NT-Xent, DINO, VICReg, Barlow Twins, CosineRegression, CombinedLoss
202
+ ├── augmentation/ # 7 transforms, compose(), MultiView
203
+ ├── nn/ # MLP, Projector, Predictor, DINOHead, pooling
204
+ ├── evaluation/ # LogRegEvaluator, KNNEvaluator, extract_embeddings
205
+ ├── training/ # DINOTrainer + callbacks
206
+ ├── data/ # DataModule (full-batch and NeighborLoader)
207
+ └── utils/ # update_ema_params, CosineDecayScheduler, CosineEMAScheduler, PositiveMiner
208
+ ```
209
+
210
+ ---
211
+
212
+ ## Key Design Choices
213
+
214
+ - **Uniform API.** Every model is `ModelClass(config: Dict, in_channels: int)`. The config dict is validated by a dedicated dataclass in `schema.py`.
215
+ - **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.
216
+ - **Protected-node augmentations.** When using `NeighborLoader` for mini-batch training, seed nodes are shielded from `node_drop` so the loss remains valid.
217
+ - **Registry pattern.** All components are registered by name (`@ENCODERS.register("gin")`), so swapping architectures is a one-line config change.
218
+ - **Composable losses.** `CombinedLoss` enables weighted mixtures of any registered losses, configurable from YAML or programmatically.
219
+
220
+ ---
221
+
222
+ ## Acknowledgments
223
+
224
+ Developed at [NECSTLab](https://necst.it), Politecnico di Milano.
@@ -0,0 +1,118 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77.0.3"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "graphssl"
7
+ version = "0.1.0"
8
+ description = "A modular Python library for Self-Supervised Learning on graphs, built on PyTorch and PyTorch Geometric."
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ authors = [
14
+ {name = "Simone Copetti", email = "copetti.simone7@gmail.com"},
15
+ ]
16
+ keywords = [
17
+ "graph-learning",
18
+ "self-supervised-learning",
19
+ "gnn",
20
+ "pytorch-geometric",
21
+ "contrastive-learning",
22
+ "graph-neural-networks",
23
+ ]
24
+ classifiers = [
25
+ "Development Status :: 3 - Alpha",
26
+ "Intended Audience :: Science/Research",
27
+ "Programming Language :: Python :: 3",
28
+ "Programming Language :: Python :: 3.10",
29
+ "Programming Language :: Python :: 3.11",
30
+ "Programming Language :: Python :: 3.12",
31
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
32
+ "Topic :: Software Development :: Libraries :: Python Modules",
33
+ ]
34
+ # Core deps: only what is always required to import graphssl.
35
+ # torch-scatter / torch-sparse are implicit via torch_geometric extras;
36
+ # users install them separately following the PyG installation guide.
37
+ dependencies = [
38
+ "torch>=2.0.0",
39
+ "torch_geometric>=2.3.0",
40
+ ]
41
+
42
+ [project.urls]
43
+ Homepage = "https://github.com/Simocop7/graphssl"
44
+ Repository = "https://github.com/Simocop7/graphssl"
45
+ "Bug Tracker" = "https://github.com/Simocop7/graphssl/issues"
46
+ Documentation = "https://github.com/Simocop7/graphssl#readme"
47
+
48
+ [project.optional-dependencies]
49
+ # Visualization: UMAP embeddings + matplotlib/seaborn plots
50
+ viz = ["umap-learn", "matplotlib", "seaborn"]
51
+ # Full: everything including AFGRL (faiss), benchmarks (ogb) and YAML configs (pyyaml)
52
+ full = ["faiss-cpu", "umap-learn", "matplotlib", "seaborn", "ogb", "pyyaml"]
53
+ # Minimal extras needed to run the benchmark example scripts
54
+ benchmark = ["ogb", "pyyaml"]
55
+ # Development: testing, building and publishing
56
+ dev = ["pytest", "pytest-cov", "build", "twine", "pyyaml", "ruff", "mypy", "pre-commit"]
57
+
58
+ [tool.setuptools.packages.find]
59
+ where = ["src"]
60
+ include = ["graphssl*"]
61
+
62
+ [tool.ruff]
63
+ line-length = 100
64
+ target-version = "py310"
65
+ src = ["src", "tests", "examples"]
66
+ # Don't reformat Python snippets embedded in prose docs (CLAUDE.md, README.md) —
67
+ # ruff format's Markdown support would touch documentation content unrelated
68
+ # to this tooling change.
69
+ extend-exclude = ["*.md"]
70
+
71
+ [tool.ruff.lint]
72
+ # E/W: pycodestyle, F: pyflakes, I: isort, UP: pyupgrade, B: bugbear, C4: comprehensions
73
+ select = ["E", "F", "W", "I", "UP", "B", "C4"]
74
+ ignore = [
75
+ "E501", # line length is handled by the formatter, not worth hard-failing on
76
+ # Typing modernization (List -> list, Optional[X] -> X | None) touches nearly
77
+ # every file for a purely cosmetic gain; deferred to a dedicated follow-up
78
+ # instead of bundling it into an unrelated tooling change. See CONTRIBUTING.md.
79
+ "UP006", "UP007", "UP035", "UP037", "UP045",
80
+ ]
81
+
82
+ [tool.ruff.lint.isort]
83
+ known-first-party = ["graphssl"]
84
+
85
+ [tool.ruff.lint.per-file-ignores]
86
+ # __init__.py files intentionally re-export names for the public API surface.
87
+ "**/__init__.py" = ["F401"]
88
+
89
+ [tool.ruff.format]
90
+ quote-style = "double"
91
+
92
+ [tool.coverage.run]
93
+ source = ["src/graphssl"]
94
+ branch = true
95
+
96
+ [tool.coverage.report]
97
+ exclude_lines = [
98
+ "pragma: no cover",
99
+ "raise NotImplementedError",
100
+ "if __name__ == .__main__.:",
101
+ "if TYPE_CHECKING:",
102
+ ]
103
+
104
+ [tool.mypy]
105
+ python_version = "3.10"
106
+ ignore_missing_imports = true
107
+ warn_unused_ignores = true
108
+ warn_redundant_casts = true
109
+ no_implicit_optional = true
110
+ # Incremental adoption: start lenient (no error on untyped defs across the whole
111
+ # codebase yet), tightened module-by-module over time. See CONTRIBUTING.md.
112
+ disallow_untyped_defs = false
113
+ check_untyped_defs = true
114
+
115
+ [[tool.mypy.overrides]]
116
+ # Third-party deps without bundled type stubs.
117
+ module = ["torch_geometric.*", "faiss.*", "sklearn.*", "umap.*", "ogb.*"]
118
+ ignore_missing_imports = true
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+