Models & adapters¶

GraphNetz follows a simple rule: a node-level encoder should run on every task without rewriting code. The five built-in architectures plug into node classification, graph classification, graph regression, and link prediction through three small adapter wrappers; GIN keeps its native graph-level pooling.

Built-in architectures¶

Model	Supported tasks	Reference
`GCN`	all four	Kipf & Welling, ICLR 2017
`GAT`	all four	Veličković et al., ICLR 2018
`GIN`	`graph_cls`, `graph_reg`	Xu et al., ICLR 2019
`GraphSAGE`	all four	Hamilton et al., NeurIPS 2017
`GraphTransformer`	all four	Shi et al., 2021
`DGI`	(utility)	Veličković et al., ICLR 2019

Note

DGI is exposed as a self-supervised pre-training utility, not a benchmark task. Its loss is its own metric — there is no held-out signal — so the runner routes unlabelled graphs through link_pred (a real edge-split AUC) instead of DGI loss.

Task type adapters¶

Three thin wrappers make a node-level encoder fluent in every task family. The benchmark dispatcher picks the right one automatically based on the chosen task_type:

class graphnetz.models._adapters.GraphLevelWrapper(*args: Any, **kwargs: Any)[source]

Bases: Module

Wrap a node-level encoder for graph-level prediction.

The encoder is expected to map a PyG Data batch to per-node features of shape [N, hidden_channels]. The wrapper adds a global mean pool over the batch index and a linear classification/regression head.

forward(data: torch_geometric.data.Data) → torch.Tensor[source]

class graphnetz.models._adapters.LinkPredWrapper(*args: Any, **kwargs: Any)[source]

Bases: Module

Wrap any node-level encoder as a link predictor with a dot-product decoder.

The wrapper exposes encode(data) returning per-node embeddings of shape [N, hidden_channels] and decode(z, edge_label_index) returning a [E] tensor of edge logits.

encode(data: torch_geometric.data.Data) → torch.Tensor[source]

static decode(z: torch.Tensor, edge_label_index: torch.Tensor) → torch.Tensor[source]

forward(data: torch_geometric.data.Data) → torch.Tensor[source]

class graphnetz.models._adapters.DGIWrapper(*args: Any, **kwargs: Any)[source]

Bases: Module

Wrap any node-level encoder as a Deep Graph Infomax model.

Mirrors the graphnetz.models.DGI interface (forward(data) returning the (pos_z, neg_z, summary) triple, plus a loss(...) helper) so the benchmark trainer does not need to special-case it.

forward(data: torch_geometric.data.Data) → tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]

loss(pos_z: torch.Tensor, neg_z: torch.Tensor, summary: torch.Tensor) → torch.Tensor[source]

You won’t usually instantiate them directly — run_benchmark wraps your encoder for you:

from graphnetz import GAT
model = GAT(in_channels=8, hidden_channels=64, out_channels=4)
# run_benchmark wraps `model` in the right adapter for the requested task type.

Custom models¶

Models declare which task types they support and the dispatcher skips incompatible (model, task) pairs. Three integration paths cover the common cases:

1. Decorator (recommended for libraries)¶

Permanent registration; the model is then visible to every run_benchmark call by name:

import torch
from graphnetz import register_model

@register_model(task_type={"node_cls", "graph_cls"})
class MyGNN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        ...

2. Class attribute (no decorator)¶

If you’d rather not depend on register_model at import time:

class MyGNN(torch.nn.Module):
    task_types = {"node_cls", "graph_cls"}
    ...

3. Inline tuple (one-shot experiments)¶

Useful for hyperparameter sweeps where each variant needs a different factory:

from graphnetz import run_benchmark

run_benchmark(
    "social",
    {
        "MyGNN-d0.3": (MyGNN, "node_cls", lambda i, h, o: MyGNN(i, h, o, dropout=0.3)),
        "MyGNN-d0.5": (MyGNN, "node_cls", lambda i, h, o: MyGNN(i, h, o, dropout=0.5)),
    },
)

Multi task factory¶

For a node-level encoder that should work across all four task types without you writing the adapter glue:

from graphnetz.benchmark import _multi_task_factory, register_model

class MyEncoder(torch.nn.Module):
    """Returns per-node embeddings of shape [N, hidden_channels]."""
    ...

_ALL_KINDS = {"node_cls", "graph_cls", "graph_reg", "link_pred"}
register_model(MyEncoder, tasks=_ALL_KINDS, factory=_multi_task_factory(MyEncoder))

Choosing an integration path¶

Scenario	Use
You’re publishing a new architecture	Decorator — clean import surface, name-based discovery
You’re benchmarking someone else’s encoder	Class attribute — no edits to upstream code beyond adding `task_types`
You’re sweeping over hyperparameters	Inline tuple — one factory per variant in the same `run_benchmark` call
You have a node-level encoder that should run everywhere	Multi-task factory — `_multi_task_factory(cls)` does the wrapping