Evaluation API

extract_embeddings

extract_embeddings(model: 'torch.nn.Module', datamodule: DataModule, encoder_source: Literal['auto', 'teacher', 'student', 'online', 'target', 'encoder'] = 'auto', device: str = 'cpu', mini_batch: bool = False, num_neighbors: Optional[List[int]] = None) -> Tuple[Tensor, Optional[Tensor]]

Extract embeddings from a model using a DataModule.

Three extraction paths:

• Graph-level (datamodule.is_graph_level=True): iterates eval_dataloader(), one embedding per graph.
• Node full-batch (default, mini_batch=False): single forward pass over the full graph stored in datamodule.data.
• Node mini-batch (mini_batch=True): NeighborLoader with global_to_local reassembly via batch.n_id. Use this for graphs too large to fit in memory for a single forward pass.

When encoder_source="auto" (default), calls model.forward() directly. This respects each model's own evaluation convention — e.g. teacher encoder for GraphDINO, online encoder for BGRL/AFGRL — without requiring the caller to know the internal architecture.

Pass an explicit encoder_source to bypass model.forward() and extract from a specific sub-encoder (useful for ablations).

Parameters:

Name	Type	Description	Default
model	'torch.nn.Module'	Any BaseSSLModel or compatible nn.Module.	required
datamodule	DataModule	DataModule exposing eval_dataloader(), .data, and neighbor_loader() (required for mini_batch=True).	required
encoder_source	Literal['auto', 'teacher', 'student', 'online', 'target', 'encoder']	"auto" calls model.forward(); one of "teacher", "student", "online", "target", "encoder" selects a specific sub-encoder directly.	'auto'
device	str	Inference device string, e.g. "cuda:0".	'cpu'
mini_batch	bool	If True, use NeighborLoader for node-level extraction. Ignored when is_graph_level=True.	False
num_neighbors	Optional[List[int]]	Neighbours per hop for NeighborLoader. Defaults to [-1, -1] (all neighbours, 2 hops). Used only when mini_batch=True.	None

Returns:

Type	Description
Tensor	(embeddings, labels) — embeddings tensor and optional labels tensor.
Optional[Tensor]	Labels are None when the dataset has no y attribute.

Source code in src/graphssl/evaluation/visualization.py

def extract_embeddings(
    model: "torch.nn.Module",
    datamodule: DataModule,
    encoder_source: Literal["auto", "teacher", "student", "online", "target", "encoder"] = "auto",
    device: str = "cpu",
    mini_batch: bool = False,
    num_neighbors: Optional[List[int]] = None,
) -> Tuple[Tensor, Optional[Tensor]]:
    """Extract embeddings from a model using a DataModule.

    Three extraction paths:

    - **Graph-level** (``datamodule.is_graph_level=True``): iterates
      ``eval_dataloader()``, one embedding per graph.
    - **Node full-batch** (default, ``mini_batch=False``): single forward pass
      over the full graph stored in ``datamodule.data``.
    - **Node mini-batch** (``mini_batch=True``): ``NeighborLoader`` with
      ``global_to_local`` reassembly via ``batch.n_id``. Use this for graphs
      too large to fit in memory for a single forward pass.

    When ``encoder_source="auto"`` (default), calls ``model.forward()``
    directly. This respects each model's own evaluation convention — e.g.
    teacher encoder for GraphDINO, online encoder for BGRL/AFGRL — without
    requiring the caller to know the internal architecture.

    Pass an explicit ``encoder_source`` to bypass ``model.forward()`` and
    extract from a specific sub-encoder (useful for ablations).

    Args:
        model: Any ``BaseSSLModel`` or compatible ``nn.Module``.
        datamodule: ``DataModule`` exposing ``eval_dataloader()``, ``.data``,
            and ``neighbor_loader()`` (required for ``mini_batch=True``).
        encoder_source: ``"auto"`` calls ``model.forward()``; one of
            ``"teacher"``, ``"student"``, ``"online"``, ``"target"``,
            ``"encoder"`` selects a specific sub-encoder directly.
        device: Inference device string, e.g. ``"cuda:0"``.
        mini_batch: If ``True``, use ``NeighborLoader`` for node-level
            extraction. Ignored when ``is_graph_level=True``.
        num_neighbors: Neighbours per hop for ``NeighborLoader``. Defaults to
            ``[-1, -1]`` (all neighbours, 2 hops). Used only when
            ``mini_batch=True``.

    Returns:
        ``(embeddings, labels)`` — embeddings tensor and optional labels tensor.
        Labels are ``None`` when the dataset has no ``y`` attribute.
    """
    _device = torch.device(device)
    model = model.to(_device)
    model.eval()

    use_model_forward = encoder_source == "auto"
    if not use_model_forward:
        _enc = _select_encoder(model, encoder_source)

    def _extract(batch) -> Tensor:
        if use_model_forward:
            return model(batch)
        return encode(_enc, batch)

    with torch.no_grad():
        if datamodule.is_graph_level:
            return _extract_graph_level(datamodule, _extract, _device)
        if mini_batch:
            return _extract_node_mini_batch(datamodule, _extract, _device, num_neighbors)
        return _extract_node_full_batch(datamodule, _extract, _device)

LogRegEvaluator

LogRegEvaluator(*, weight_decay: Union[float, Sequence[float]] = DEFAULT_WEIGHT_DECAYS, max_iter: int = 5000, standardize: bool = True, multilabel: bool = False)

L2-regularised logistic regression on frozen embeddings.

The classifier minimises mean cross-entropy + weight_decay / 2 * ||W||² on the training rows, full-batch with L-BFGS, starting from zero: no random initialisation and no learning rate, so the same embeddings always give the same result. When several weight_decay values are given, the one with the best validation metric is used (the strongest on ties) and the test metric is reported for it. The values are fitted from the strongest to the weakest, each starting from the previous solution.

For multilabel tasks (e.g. ogbg-molpcba), uses binary cross-entropy and reports average precision score instead of accuracy. Labels are expected as float tensors possibly containing NaN for missing values.

Parameters:

Name	Type	Description	Default
weight_decay	Union[float, Sequence[float]]	L2 strength, or the values to select among on the validation split.	DEFAULT_WEIGHT_DECAYS
max_iter	int	Budget of L-BFGS iterations and of loss evaluations for each fit. converged in the result is False when the selected fit used it up.	5000
standardize	bool	Rescale every feature to zero mean and unit variance, with the statistics of the training rows only.	True
multilabel	bool	If True, treats the task as multilabel classification.	False

evaluate() returns val_acc / test_acc (val_ap / test_ap when multilabel), the selected weight_decay and converged.

Source code in src/graphssl/evaluation/linear_probe.py

def __init__(
    self,
    *,
    weight_decay: Union[float, Sequence[float]] = DEFAULT_WEIGHT_DECAYS,
    max_iter: int = 5000,
    standardize: bool = True,
    multilabel: bool = False,
):
    values = [weight_decay] if isinstance(weight_decay, (int, float)) else list(weight_decay)
    if not values or any(v < 0 for v in values):
        raise ValueError(f"weight_decay must be one or more values >= 0, got {weight_decay!r}")
    if max_iter <= 0:
        raise ValueError(f"max_iter must be > 0, got {max_iter}")
    self.weight_decays: Tuple[float, ...] = tuple(
        sorted({float(v) for v in values}, reverse=True)
    )
    self.max_iter = max_iter
    self.standardize = standardize
    self.multilabel = multilabel

RidgeEvaluator

RidgeEvaluator(*, weight_decay: Union[float, Sequence[float]] = DEFAULT_RIDGE_WEIGHT_DECAYS, standardize: bool = True)

Ridge regression on frozen embeddings, for node- or graph-level regression targets.

Fits mean squared error + weight_decay * ||W||² on the training rows in closed form: no optimiser, no random initialisation, so the same embeddings always give the same result. When several weight_decay values are given, the one with the lowest validation MAE is used (the strongest on ties) and the test error is reported for it.

Parameters:

Name	Type	Description	Default
weight_decay	Union[float, Sequence[float]]	L2 strength, or the values to select among on the validation split.	DEFAULT_RIDGE_WEIGHT_DECAYS
standardize	bool	Rescale every feature to zero mean and unit variance, with the statistics of the training rows only.	True

evaluate() returns val_mae / test_mae, val_rmse / test_rmse and the selected weight_decay. Targets are [N] or [N, T] (errors averaged over the T targets) and must not contain NaN.

Source code in src/graphssl/evaluation/regression_probe.py

def __init__(
    self,
    *,
    weight_decay: Union[float, Sequence[float]] = DEFAULT_RIDGE_WEIGHT_DECAYS,
    standardize: bool = True,
):
    values = [weight_decay] if isinstance(weight_decay, (int, float)) else list(weight_decay)
    if not values or any(v < 0 for v in values):
        raise ValueError(f"weight_decay must be one or more values >= 0, got {weight_decay!r}")
    self.weight_decays: Tuple[float, ...] = tuple(
        sorted({float(v) for v in values}, reverse=True)
    )
    self.standardize = standardize

KNNEvaluator

KNNEvaluator(k: int = 20, chunk_size: int = 4096)

Weighted k-nearest-neighbour classifier on frozen embeddings.

Uses cosine similarity to find neighbours. Labels are assigned by majority vote among the k nearest training nodes.

Parameters:

Name	Type	Description	Default
k	int	Number of nearest neighbours.	20
chunk_size	int	Rows of the evaluated split scored at a time. The full [|split|, |train|] similarity matrix is never built (it would take ~17 GB for the ogbn-arxiv test split); the predictions are identical.	4096

Source code in src/graphssl/evaluation/knn.py

def __init__(self, k: int = 20, chunk_size: int = 4096):
    if chunk_size <= 0:
        raise ValueError(f"chunk_size must be > 0, got {chunk_size}")
    self.k = k
    self.chunk_size = chunk_size

effective_rank

effective_rank(embeddings: Tensor, eps: float = 1e-12) -> float

Effective rank of an embedding matrix (Roy & Vetterli, 2007).

exp(H(p)), where p are the singular values of the column-centered embeddings [N, D] normalized to sum to 1. It ranges from 1 (all variance along a single direction) to min(N, D) (variance spread evenly), and is 0 when the embeddings have no variance at all. A value far below D signals dimensional collapse, a common failure mode of non-contrastive SSL objectives.

Source code in src/graphssl/evaluation/diagnostics.py

@torch.no_grad()
def effective_rank(embeddings: Tensor, eps: float = 1e-12) -> float:
    """Effective rank of an embedding matrix (Roy & Vetterli, 2007).

    ``exp(H(p))``, where ``p`` are the singular values of the column-centered
    ``embeddings`` [N, D] normalized to sum to 1. It ranges from 1 (all variance along a
    single direction) to ``min(N, D)`` (variance spread evenly), and is 0 when the
    embeddings have no variance at all. A value far below ``D`` signals dimensional
    collapse, a common failure mode of non-contrastive SSL objectives.
    """
    z = embeddings.detach().float()
    s = torch.linalg.svdvals(z - z.mean(dim=0, keepdim=True))
    total = s.sum()
    if total <= eps:
        return 0.0
    p = s / total
    entropy = -(p * p.clamp_min(eps).log()).sum()
    return float(entropy.exp())