gnn-lab

A clean reference implementation of GNNs across two task families:

• Graph classification — MUTAG, PROTEINS
• Graph regression — ZINC, ESOL

Five architectures: GIN, GINE, GCN, GraphSAGE, GAT. Each command runs all models and classical ML baselines in one go, producing a summary table.

---

Takeaway

Practical things to keep in mind for future graph tasks:

Always run a classical baseline first. WL-SVM and ECFP-RF are cheap to fit and surprisingly strong. If your GNN can't beat them, you likely have a data-size or overfitting problem worth diagnosing before scaling up.

Check whether your dataset has edge features — and pick your architecture accordingly. If edge attributes carry signal (bond type, distance, etc.), use GINE or another edge-aware conv. GIN ignores them entirely and you'll leave performance on the table.

On small datasets (< ~1k graphs), trust cross-validated accuracy over single splits. MUTAG has 188 graphs; a single 80/20 split is essentially noise. Use k-fold CV and report mean ± std — if the std is large, the ranking between models is unreliable.

GAT needs care. It collapsed on PROTEINS despite being a reasonable choice on paper. Attention adds parameters and variance; tune the number of heads and consider whether your dataset is large enough to support it.

For regression, pooling choice matters more than for classification. Sum pooling (GIN/GINE) outperforms mean pooling (GCN, SAGE, GAT) when the target correlates with graph size or atom count — common in molecular property prediction.

---

Installation

Requires Python ≥ 3.13 and uv.

git clone https://github.com/th-tsai/gnn-lab
cd gnn-lab
uv sync

---

Quick Start

uv run main.py train-graph-cls   # classification: MUTAG + PROTEINS
uv run main.py train-graph-reg   # regression: ZINC + ESOL

Each command reads its config, runs every baseline and GNN model on every dataset, then prints a summary table.

---

Configuration

Edit config/train-graph-cls.yaml or config/train-graph-reg.yaml to change datasets, models, or hyperparameters — no CLI flags needed.

config/train-graph-cls.yaml

datasets: [MUTAG, PROTEINS]
models: [gin, gcn, sage, gat]
baselines: [wl-svm, ecfp-rf]

seed: 42
folds: 10
epochs: 100
hidden_channels: 64
num_layers: 5
dropout: 0.5
lr: 0.01
batch_size: 32

config/train-graph-reg.yaml

datasets: [zinc, esol]
models: [gin, gine, gcn, sage, gat]
baselines: [ecfp-rf]

seed: 42
epochs: 300
hidden_channels: 64
num_layers: 5
dropout: 0.0
lr: 0.001
batch_size: 128

---

Datasets

Graph classification

Dataset	Graphs	Avg nodes	Classes	Node features
MUTAG	188	17.9	2	7 (atom type, one-hot)
PROTEINS	1,113	39.1	2	3 (continuous)

Graph regression

Dataset	Graphs	Target	Metric
ZINC (subset)	12,000	Penalized logP	MAE
ESOL	1,128	Log aqueous solubility	MAE

---

Models

Model	Conv	Pool	Edge features	Notes
GIN	GINConv + BN	sum	no	Max expressiveness (Xu et al. 2019)
GINE	GINEConv	sum	yes	GIN + edge features; used for ZINC, ESOL
GCN	GCNConv + BN	mean	no	Kipf & Welling 2017
GraphSAGE	SAGEConv + BN	mean	no	Hamilton et al. 2017
GAT	GATConv + BN	mean	no	Veličković et al. 2018

GINE is automatically skipped for datasets without edge features.

---

Classical Baselines

Method	Task	Description
wl-svm	cls	WL subtree kernel features + LinearSVC, k-fold CV
ecfp-rf	cls, reg	Morgan fingerprint (RDKit, radius=2, 2048 bits) + RandomForest

ecfp-rf builds fingerprints from data.smiles (MoleculeNet datasets like ESOL) or by reconstructing an RDKit Mol from ZINC's atom_type / bond_type integer codes. Auto-skips on non-molecular datasets (e.g., PROTEINS, MUTAG via TUDataset).

---

Results

Graph classification (10-fold CV accuracy %)

Model	MUTAG	PROTEINS
WL-SVM	84.6 ± 6.1	73.1 ± 3.4
GIN	84.0 ± 8.5	67.8 ± 7.2
GCN	76.6 ± 7.5	72.8 ± 4.9
GraphSAGE	78.7 ± 5.5	73.0 ± 4.9
GAT	70.2 ± 12.2	40.4 ± 0.2

Graph regression (test MAE)

Model	ZINC	ESOL
ECFP-RF	0.7578	0.8075
GIN	0.3648	0.3846
GINE	0.2697	0.4565
GCN	0.7510	0.4666
GraphSAGE	0.5843	0.4624
GAT	0.6386	0.4159

---

Repository Structure

gnn-lab/
├── main.py
├── pyproject.toml
├── config/
│   ├── train-graph-cls.yaml
│   └── train-graph-reg.yaml
├── src/gnn_lab/
│   ├── graph_cls.py      # train-graph-cls entry point
│   ├── graph_reg.py      # train-graph-reg entry point
│   ├── baselines.py      # WL-SVM and ECFP-RF baselines
│   ├── model.py          # GIN, GINE, GCN, GraphSAGE, GAT
│   ├── train.py          # Training loops
│   └── dataset.py        # Dataset loaders
└── tests/
    └── test_smoke.py     # Model shape + one-step training smoke tests

---

References

• GIN / GINE: Xu et al. (2019) How Powerful are Graph Neural Networks?
• GCN: Kipf & Welling (2017) Semi-Supervised Classification with Graph Convolutional Networks
• GraphSAGE: Hamilton et al. (2017) Inductive Representation Learning on Large Graphs
• GAT: Veličković et al. (2018) Graph Attention Networks
• ESOL: Delaney (2004) ESOL: Estimating Aqueous Solubility Directly from Molecular Structure