graphnetz.training

training

Lightweight training loops shared by the example notebooks.

Each function returns a plain dict of per-epoch metrics, ready to feed into plot_history.

All trainers accept device='auto' (the default), which dispatches to CUDA when available, then Apple-silicon MPS, then CPU. Pass an explicit torch.device or string to pin placement.

Functions:

Name	Description
train_node_classification	Train a node classifier with Planetoid-style train/val/test masks.
train_graph_classification	Train a graph-level classifier.
train_graph_regression	Train a graph-level regressor (MSE loss, MAE on val).
train_node_degree_regression	Self-supervised node-level regression: predict log node degree.
train_dgi	Train a Deep Graph Infomax model (unsupervised).
train_link_prediction	Train a link predictor with binary cross-entropy on RandomLinkSplit.
train_relational_link_prediction	Train a relational link predictor (DistMult) on knowledge graph triples.

train_node_classification

train_node_classification(
    model: Module,
    data: Data,
    epochs: int = 100,
    lr: float = 0.01,
    weight_decay: float = 0.0005,
    verbose: bool = False,
    device: device | str | None = "auto",
) -> dict[str, list[float]]

Train a node classifier with Planetoid-style train/val/test masks.

Source code in src/graphnetz/training.py

def train_node_classification(
    model: torch.nn.Module,
    data: Data,
    epochs: int = 100,
    lr: float = 1e-2,
    weight_decay: float = 5e-4,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
) -> dict[str, list[float]]:
    """Train a node classifier with Planetoid-style train/val/test masks."""
    dev = _resolve_device(device)
    model = model.to(dev)
    data = data.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
    history: dict[str, list[float]] = {"train_loss": [], "val_acc": [], "test_acc": []}
    train_mask = _select_split_mask(data.train_mask)
    val_mask = _select_split_mask(data.val_mask)
    test_mask = _select_split_mask(data.test_mask)
    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        opt.zero_grad()
        out = model(data)
        loss = F.cross_entropy(out[train_mask], data.y[train_mask])
        loss.backward()
        opt.step()

        model.eval()
        with torch.no_grad():
            pred = model(data).argmax(dim=1)
            val_acc = (pred[val_mask] == data.y[val_mask]).float().mean().item()
            test_acc = (pred[test_mask] == data.y[test_mask]).float().mean().item()
        history["train_loss"].append(loss.item())
        history["val_acc"].append(val_acc)
        history["test_acc"].append(test_acc)
        if verbose:
            iterator.set_postfix(
                loss=f"{loss.item():.4f}",
                val=f"{val_acc:.4f}",
                test=f"{test_acc:.4f}",
            )
    return history

train_graph_classification

train_graph_classification(
    model: Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    epochs: int = 30,
    lr: float = 0.001,
    verbose: bool = False,
    device: device | str | None = "auto",
    test_loader: DataLoader | None = None,
) -> dict[str, list[float]]

Train a graph-level classifier.

Handles single-label and multi-label datasets transparently: when batch.y is shaped [B, C] with float dtype (e.g. LRGB Peptides-func, OGB molhiv variants), the loss switches to binary cross-entropy with logits and the reported metric is the average correctly-classified label fraction.

With a test_loader the history also carries test_acc, a held-out series distinct from the val_acc one, which is what makes epoch_selection="best_val" and hyperparameter search legitimate.

Source code in src/graphnetz/training.py

def train_graph_classification(
    model: torch.nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    epochs: int = 30,
    lr: float = 1e-3,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
    test_loader: DataLoader | None = None,
) -> dict[str, list[float]]:
    """Train a graph-level classifier.

    Handles single-label and multi-label datasets transparently:
    when ``batch.y`` is shaped ``[B, C]`` with float dtype (e.g.\
    LRGB ``Peptides-func``, OGB molhiv variants), the loss switches to
    binary cross-entropy with logits and the reported metric is the
    average correctly-classified label fraction.

    With a ``test_loader`` the history also carries ``test_acc``, a held-out
    series distinct from the ``val_acc`` one, which is what makes
    ``epoch_selection="best_val"`` and hyperparameter search legitimate.
    """
    dev = _resolve_device(device)
    model = model.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history: dict[str, list[float]] = {"train_loss": [], "val_acc": []}
    if test_loader is not None:
        history["test_acc"] = []

    # Sniff one batch to decide single- vs. multi-label.
    sniff_batch = next(iter(train_loader))
    if sniff_batch.y is None:
        msg = "train_graph_classification requires graph-level `y` on each batch"
        raise ValueError(msg)
    multi_label = sniff_batch.y.dim() == 2 and sniff_batch.y.size(1) > 1

    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        total = 0.0
        n = 0
        for batch in train_loader:
            batch = batch.to(dev)
            opt.zero_grad()
            out = model(batch)
            if multi_label:
                # Multi-label targets may be partially unlabelled (NaN, e.g.
                # ogbg-molpcba); only the observed entries enter the loss.
                target = batch.y.float()
                observed = ~torch.isnan(target)
                loss = F.binary_cross_entropy_with_logits(out[observed], target[observed])
            else:
                loss = F.cross_entropy(out, batch.y.view(-1))
            loss.backward()
            opt.step()
            total += loss.item() * batch.num_graphs
            n += batch.num_graphs

        model.eval()

        def _accuracy(loader: DataLoader) -> float:
            correct = 0.0
            m = 0
            with torch.no_grad():
                for batch in loader:
                    batch = batch.to(dev)
                    if multi_label:
                        pred = (model(batch) > 0).float()
                        target = batch.y.float()
                        observed = ~torch.isnan(target)
                        if observed.any():
                            hit = (pred[observed] == target[observed]).float().mean().item()
                            correct += hit * batch.num_graphs
                    else:
                        pred = model(batch).argmax(dim=1)
                        correct += (pred == batch.y.view(-1)).sum().item()
                    m += batch.num_graphs
            return correct / max(m, 1)

        train_loss = total / max(n, 1)
        val_acc = _accuracy(val_loader)
        history["train_loss"].append(train_loss)
        history["val_acc"].append(val_acc)
        if test_loader is not None:
            history["test_acc"].append(_accuracy(test_loader))
        if verbose:
            iterator.set_postfix(loss=f"{train_loss:.4f}", val=f"{val_acc:.4f}")
    return history

train_graph_regression

train_graph_regression(
    model: Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    epochs: int = 30,
    lr: float = 0.001,
    verbose: bool = False,
    device: device | str | None = "auto",
    test_loader: DataLoader | None = None,
) -> dict[str, list[float]]

Train a graph-level regressor (MSE loss, MAE on val).

With a test_loader the history also carries test_mae on a held-out split distinct from the validation one.

Source code in src/graphnetz/training.py

def train_graph_regression(
    model: torch.nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    epochs: int = 30,
    lr: float = 1e-3,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
    test_loader: DataLoader | None = None,
) -> dict[str, list[float]]:
    """Train a graph-level regressor (MSE loss, MAE on val).

    With a ``test_loader`` the history also carries ``test_mae`` on a
    held-out split distinct from the validation one.
    """
    dev = _resolve_device(device)
    model = model.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history: dict[str, list[float]] = {"train_loss": [], "val_mae": []}
    if test_loader is not None:
        history["test_mae"] = []
    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        total = 0.0
        n = 0
        for batch in train_loader:
            batch = batch.to(dev)
            opt.zero_grad()
            out = model(batch).view(-1)
            loss = F.mse_loss(out, batch.y.float().view(-1))
            loss.backward()
            opt.step()
            total += loss.item() * batch.num_graphs
            n += batch.num_graphs

        model.eval()

        def _mae(loader: DataLoader) -> float:
            mae = 0.0
            m = 0
            with torch.no_grad():
                for batch in loader:
                    batch = batch.to(dev)
                    out = model(batch).view(-1)
                    mae += (out - batch.y.float().view(-1)).abs().sum().item()
                    m += batch.num_graphs
            return mae / max(m, 1)

        train_loss = total / max(n, 1)
        val_mae = _mae(val_loader)
        history["train_loss"].append(train_loss)
        history["val_mae"].append(val_mae)
        if test_loader is not None:
            history["test_mae"].append(_mae(test_loader))
        if verbose:
            iterator.set_postfix(loss=f"{train_loss:.4f}", mae=f"{val_mae:.4f}")
    return history

train_node_degree_regression

train_node_degree_regression(
    model: Module,
    data: Data,
    epochs: int = 100,
    lr: float = 0.01,
    verbose: bool = False,
    device: device | str | None = "auto",
) -> dict[str, list[float]]

Self-supervised node-level regression: predict log node degree.

Source code in src/graphnetz/training.py

def train_node_degree_regression(
    model: torch.nn.Module,
    data: Data,
    epochs: int = 100,
    lr: float = 1e-2,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
) -> dict[str, list[float]]:
    """Self-supervised node-level regression: predict log node degree."""
    dev = _resolve_device(device)
    model = model.to(dev)
    data = data.to(dev)
    target = torch.log1p(degree(data.edge_index[0], num_nodes=data.num_nodes).float())
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history: dict[str, list[float]] = {"train_loss": [], "val_mae": []}
    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        opt.zero_grad()
        out = model(data).view(-1)
        loss = F.mse_loss(out, target)
        loss.backward()
        opt.step()
        model.eval()
        with torch.no_grad():
            eval_out = model(data).view(-1)
            mae = (eval_out - target).abs().mean().item()
        history["train_loss"].append(loss.item())
        history["val_mae"].append(mae)
        if verbose:
            iterator.set_postfix(loss=f"{loss.item():.4f}", mae=f"{mae:.4f}")
    return history

train_dgi

train_dgi(
    model: _DGILike,
    data: Data,
    epochs: int = 100,
    lr: float = 0.001,
    verbose: bool = False,
    device: device | str | None = "auto",
) -> dict[str, list[float]]

Train a Deep Graph Infomax model (unsupervised).

Source code in src/graphnetz/training.py

def train_dgi(
    model: _DGILike,
    data: Data,
    epochs: int = 100,
    lr: float = 1e-3,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
) -> dict[str, list[float]]:
    """Train a Deep Graph Infomax model (unsupervised)."""
    dev = _resolve_device(device)
    if isinstance(model, torch.nn.Module):
        model = model.to(dev)
    data = data.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history: dict[str, list[float]] = {"dgi_loss": []}
    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        opt.zero_grad()
        pos_z, neg_z, summary = model(data)
        loss = model.loss(pos_z, neg_z, summary)
        loss.backward()
        opt.step()
        history["dgi_loss"].append(loss.item())
        if verbose:
            iterator.set_postfix(loss=f"{loss.item():.4f}")
    return history

train_link_prediction

train_link_prediction(
    model: _LinkPredLike,
    train_data: Data,
    val_data: Data,
    test_data: Data,
    epochs: int = 100,
    lr: float = 0.01,
    verbose: bool = False,
    device: device | str | None = "auto",
) -> dict[str, list[float]]

Train a link predictor with binary cross-entropy on RandomLinkSplit.

The model is expected to expose encode(data) returning per-node embeddings and decode(z, edge_label_index) returning per-edge scores (see LinkPredWrapper).

Source code in src/graphnetz/training.py

def train_link_prediction(
    model: _LinkPredLike,
    train_data: Data,
    val_data: Data,
    test_data: Data,
    epochs: int = 100,
    lr: float = 1e-2,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
) -> dict[str, list[float]]:
    """Train a link predictor with binary cross-entropy on RandomLinkSplit.

    The model is expected to expose ``encode(data)`` returning per-node
    embeddings and ``decode(z, edge_label_index)`` returning per-edge scores
    (see [`LinkPredWrapper`][graphnetz.models._adapters.LinkPredWrapper]).
    """
    dev = _resolve_device(device)
    if isinstance(model, torch.nn.Module):
        model = model.to(dev)
    train_data = train_data.to(dev)
    val_data = val_data.to(dev)
    test_data = test_data.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history: dict[str, list[float]] = {"train_loss": [], "val_auc": [], "test_auc": []}
    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        opt.zero_grad()
        z = model.encode(train_data)
        logits = model.decode(z, train_data.edge_label_index)
        loss = F.binary_cross_entropy_with_logits(logits, train_data.edge_label.float())
        loss.backward()
        opt.step()

        model.eval()
        with torch.no_grad():
            z_eval = model.encode(train_data)
            val_logits = model.decode(z_eval, val_data.edge_label_index).cpu()
            test_logits = model.decode(z_eval, test_data.edge_label_index).cpu()
            val_auc = _safe_auc(val_data.edge_label.cpu(), val_logits)
            test_auc = _safe_auc(test_data.edge_label.cpu(), test_logits)
        history["train_loss"].append(float(loss.item()))
        history["val_auc"].append(val_auc)
        history["test_auc"].append(test_auc)
        if verbose:
            iterator.set_postfix(
                loss=f"{loss.item():.4f}",
                val=f"{val_auc:.4f}",
                test=f"{test_auc:.4f}",
            )
    return history

train_relational_link_prediction

train_relational_link_prediction(
    model: _RelationalLinkPredLike,
    train_data: Data,
    val_data: Data,
    test_data: Data,
    epochs: int = 100,
    lr: float = 0.01,
    verbose: bool = False,
    device: device | str | None = "auto",
) -> dict[str, list[float]]

Train a relational link predictor (DistMult) on knowledge graph triples.

The model is expected to expose encode(data) returning per-node embeddings and decode(z, edge_index, edge_type) returning per-edge scores (see graphnetz.models._adapters.RelationalLinkPredWrapper).

Source code in src/graphnetz/training.py

def train_relational_link_prediction(
    model: _RelationalLinkPredLike,
    train_data: Data,
    val_data: Data,
    test_data: Data,
    epochs: int = 100,
    lr: float = 1e-2,
    verbose: bool = False,
    device: torch.device | str | None = "auto",
) -> dict[str, list[float]]:
    """Train a relational link predictor (DistMult) on knowledge graph triples.

    The model is expected to expose ``encode(data)`` returning per-node
    embeddings and ``decode(z, edge_index, edge_type)`` returning per-edge
    scores (see `graphnetz.models._adapters.RelationalLinkPredWrapper`).
    """
    dev = _resolve_device(device)
    if isinstance(model, torch.nn.Module):
        model = model.to(dev)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    history: dict[str, list[float]] = {"train_loss": [], "val_auc": [], "test_auc": []}

    # Ensure node features exist (fabricated on CPU first, then moved to
    # ``dev`` so the encoder, eye-features, and edge tensors all share
    # one device.)
    train_data = _ensure_node_features(train_data).to(dev)
    val_data = _ensure_node_features(val_data).to(dev)
    test_data = _ensure_node_features(test_data).to(dev)

    pos_edge_index = train_data.edge_index
    pos_edge_type = train_data.edge_type
    pos_labels = torch.ones(pos_edge_index.size(1), device=pos_edge_index.device)
    neg_labels = torch.zeros(pos_edge_index.size(1), device=pos_edge_index.device)
    train_labels = torch.cat([pos_labels, neg_labels])
    positive_set = {
        (int(h), int(r), int(t))
        for h, r, t in zip(pos_edge_index[0].tolist(), pos_edge_type.tolist(), pos_edge_index[1].tolist(), strict=False)
    }

    iterator = tqdm(range(epochs), desc="Epochs", leave=False, disable=not verbose)
    for _ in iterator:
        model.train()
        # Resample negatives every epoch so the model sees fresh corruptions
        # (otherwise it overfits to a fixed negative set). Filter out
        # corruptions that collide with real training positives.
        neg_edge_index, neg_edge_type = _sample_negative_triples(
            pos_edge_index, pos_edge_type, train_data.num_nodes, positive_set=positive_set
        )
        train_edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=1)
        train_edge_type = torch.cat([pos_edge_type, neg_edge_type])

        opt.zero_grad()
        z = model.encode(train_data)
        logits = model.decode(z, train_edge_index, train_edge_type)
        loss = F.binary_cross_entropy_with_logits(logits, train_labels)
        loss.backward()
        opt.step()

        model.eval()
        with torch.no_grad():
            z_eval = model.encode(train_data)
            val_auc = _relational_eval_auc(model, z_eval, val_data, train_data.num_nodes)
            test_auc = _relational_eval_auc(model, z_eval, test_data, train_data.num_nodes)
        history["train_loss"].append(float(loss.item()))
        history["val_auc"].append(val_auc)
        history["test_auc"].append(test_auc)
        if verbose:
            iterator.set_postfix(
                loss=f"{loss.item():.4f}",
                val=f"{val_auc:.4f}",
                test=f"{test_auc:.4f}",
            )
    return history