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.
- graphssl-0.1.0/LICENSE +21 -0
- graphssl-0.1.0/PKG-INFO +273 -0
- graphssl-0.1.0/README.md +224 -0
- graphssl-0.1.0/pyproject.toml +118 -0
- graphssl-0.1.0/setup.cfg +4 -0
- graphssl-0.1.0/src/graphssl/__init__.py +63 -0
- graphssl-0.1.0/src/graphssl/augmentation/__init__.py +11 -0
- graphssl-0.1.0/src/graphssl/augmentation/compose.py +60 -0
- graphssl-0.1.0/src/graphssl/augmentation/functional.py +127 -0
- graphssl-0.1.0/src/graphssl/augmentation/transforms.py +84 -0
- graphssl-0.1.0/src/graphssl/config/__init__.py +14 -0
- graphssl-0.1.0/src/graphssl/config/load.py +114 -0
- graphssl-0.1.0/src/graphssl/config/schema.py +413 -0
- graphssl-0.1.0/src/graphssl/core/__init__.py +14 -0
- graphssl-0.1.0/src/graphssl/core/augmentation.py +16 -0
- graphssl-0.1.0/src/graphssl/core/callback.py +44 -0
- graphssl-0.1.0/src/graphssl/core/encoder.py +18 -0
- graphssl-0.1.0/src/graphssl/core/model.py +58 -0
- graphssl-0.1.0/src/graphssl/core/registry.py +47 -0
- graphssl-0.1.0/src/graphssl/data/__init__.py +1 -0
- graphssl-0.1.0/src/graphssl/data/datamodule.py +100 -0
- graphssl-0.1.0/src/graphssl/encoders/__init__.py +3 -0
- graphssl-0.1.0/src/graphssl/encoders/gcn.py +104 -0
- graphssl-0.1.0/src/graphssl/encoders/gin.py +181 -0
- graphssl-0.1.0/src/graphssl/encoders/transformer.py +206 -0
- graphssl-0.1.0/src/graphssl/evaluation/__init__.py +3 -0
- graphssl-0.1.0/src/graphssl/evaluation/knn.py +52 -0
- graphssl-0.1.0/src/graphssl/evaluation/linear_probe.py +116 -0
- graphssl-0.1.0/src/graphssl/evaluation/visualization.py +206 -0
- graphssl-0.1.0/src/graphssl/losses/__init__.py +6 -0
- graphssl-0.1.0/src/graphssl/losses/barlow.py +35 -0
- graphssl-0.1.0/src/graphssl/losses/combined.py +67 -0
- graphssl-0.1.0/src/graphssl/losses/dino.py +38 -0
- graphssl-0.1.0/src/graphssl/losses/nt_xent.py +40 -0
- graphssl-0.1.0/src/graphssl/losses/regression.py +44 -0
- graphssl-0.1.0/src/graphssl/losses/vicreg.py +56 -0
- graphssl-0.1.0/src/graphssl/models/__init__.py +14 -0
- graphssl-0.1.0/src/graphssl/models/afgrl.py +133 -0
- graphssl-0.1.0/src/graphssl/models/barlow_twins.py +68 -0
- graphssl-0.1.0/src/graphssl/models/bgrl.py +110 -0
- graphssl-0.1.0/src/graphssl/models/dgi.py +85 -0
- graphssl-0.1.0/src/graphssl/models/graphcl.py +68 -0
- graphssl-0.1.0/src/graphssl/models/graphdino.py +182 -0
- graphssl-0.1.0/src/graphssl/models/supervised.py +48 -0
- graphssl-0.1.0/src/graphssl/models/vicreg.py +72 -0
- graphssl-0.1.0/src/graphssl/nn/__init__.py +4 -0
- graphssl-0.1.0/src/graphssl/nn/dino_head.py +99 -0
- graphssl-0.1.0/src/graphssl/nn/mlp.py +52 -0
- graphssl-0.1.0/src/graphssl/nn/norm.py +24 -0
- graphssl-0.1.0/src/graphssl/nn/pooling.py +18 -0
- graphssl-0.1.0/src/graphssl/registry/__init__.py +1 -0
- graphssl-0.1.0/src/graphssl/registry/registry.py +9 -0
- graphssl-0.1.0/src/graphssl/training/__init__.py +2 -0
- graphssl-0.1.0/src/graphssl/training/callbacks.py +138 -0
- graphssl-0.1.0/src/graphssl/training/trainer.py +108 -0
- graphssl-0.1.0/src/graphssl/utils/__init__.py +3 -0
- graphssl-0.1.0/src/graphssl/utils/ema.py +14 -0
- graphssl-0.1.0/src/graphssl/utils/positive_miner.py +119 -0
- graphssl-0.1.0/src/graphssl/utils/schedulers.py +55 -0
- graphssl-0.1.0/src/graphssl.egg-info/PKG-INFO +273 -0
- graphssl-0.1.0/src/graphssl.egg-info/SOURCES.txt +72 -0
- graphssl-0.1.0/src/graphssl.egg-info/dependency_links.txt +1 -0
- graphssl-0.1.0/src/graphssl.egg-info/requires.txt +29 -0
- graphssl-0.1.0/src/graphssl.egg-info/top_level.txt +1 -0
- graphssl-0.1.0/tests/test_afgrl.py +129 -0
- graphssl-0.1.0/tests/test_barlow_twins.py +108 -0
- graphssl-0.1.0/tests/test_bgrl.py +141 -0
- graphssl-0.1.0/tests/test_dgi.py +105 -0
- graphssl-0.1.0/tests/test_evaluation.py +69 -0
- graphssl-0.1.0/tests/test_graphcl.py +97 -0
- graphssl-0.1.0/tests/test_graphdino.py +257 -0
- graphssl-0.1.0/tests/test_new_features.py +255 -0
- graphssl-0.1.0/tests/test_supervised.py +107 -0
- 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.
|
graphssl-0.1.0/PKG-INFO
ADDED
|
@@ -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
|
+
[](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml)
|
|
53
|
+
[](https://codecov.io/gh/Simocop7/graphssl)
|
|
54
|
+
[](https://github.com/astral-sh/ruff)
|
|
55
|
+
[](pyproject.toml)
|
|
56
|
+
[](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.
|
graphssl-0.1.0/README.md
ADDED
|
@@ -0,0 +1,224 @@
|
|
|
1
|
+
# GraphSSL
|
|
2
|
+
|
|
3
|
+
[](https://github.com/Simocop7/graphssl/actions/workflows/tests.yml)
|
|
4
|
+
[](https://codecov.io/gh/Simocop7/graphssl)
|
|
5
|
+
[](https://github.com/astral-sh/ruff)
|
|
6
|
+
[](pyproject.toml)
|
|
7
|
+
[](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
|
graphssl-0.1.0/setup.cfg
ADDED