graphssl 0.1.0__py3-none-any.whl

graphssl/__init__.py

@@ -0,0 +1,63 @@
+"""GraphSSL: a modular Graph Machine Learning library in pure PyTorch."""
+
+# ── Core abstractions ─────────────────────────────────────────────────────────
+# ── Augmentation ──────────────────────────────────────────────────────────────
+from graphssl.augmentation import MultiView, compose
+from graphssl.augmentation.transforms import (
+    EdgeAdd,
+    EdgeDrop,
+    FeatMask,
+    FeatNoise,
+    FeatShuffle,
+    Subgraph,
+)
+from graphssl.core.augmentation import BaseAugmentation
+from graphssl.core.callback import Callback
+from graphssl.core.encoder import BaseEncoder
+from graphssl.core.model import BaseModel, BaseSSLModel
+from graphssl.core.registry import Registry
+
+# ── Data ──────────────────────────────────────────────────────────────────────
+from graphssl.data import DataModule
+
+# ── Encoders ──────────────────────────────────────────────────────────────────
+from graphssl.encoders import GCNEncoder, GINEncoder, TransformerEncoder
+
+# ── Evaluation ────────────────────────────────────────────────────────────────
+from graphssl.evaluation import KNNEvaluator, LogRegEvaluator, extract_embeddings
+
+# ── Losses ────────────────────────────────────────────────────────────────────
+from graphssl.losses import (
+    BarlowTwinsLoss,
+    CombinedLoss,
+    CosineRegressionLoss,
+    DINOLoss,
+    NTXentLoss,
+    VICRegLoss,
+)
+
+# ── Models ────────────────────────────────────────────────────────────────────
+from graphssl.models import (
+    AFGRL,
+    BGRL,
+    DGI,
+    BarlowTwins,
+    GraphCL,
+    GraphDINO,
+    Supervised,
+    VICReg,
+)
+
+# ── Neural network building blocks ────────────────────────────────────────────
+from graphssl.nn import MLP, DINOHead, Predictor, Projector, pool_graph_embeddings
+
+# ── Training ──────────────────────────────────────────────────────────────────
+from graphssl.training import (
+    DINOTrainer,
+    EmbeddingLoggerCallback,
+    LinearEvalCallback,
+    VisualizationCallback,
+)
+
+# ── Utilities ─────────────────────────────────────────────────────────────────
+from graphssl.utils import CosineDecayScheduler, CosineEMAScheduler, update_ema_params

graphssl/augmentation/__init__.py

@@ -0,0 +1,11 @@
+from . import functional
+from .compose import MultiView, compose
+from .transforms import (
+    EdgeAdd,
+    EdgeDrop,
+    FeatMask,
+    FeatNoise,
+    FeatShuffle,
+    NodeDrop,
+    Subgraph,
+)

graphssl/augmentation/compose.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+import torch
+from torch_geometric.data import Data
+
+from graphssl.registry import AUGMENTS
+
+
+def compose(data: Data, augments: list, protected_nodes: Optional[torch.Tensor] = None) -> Data:
+    """Apply a sequence of registry-based augmentations to a Data object.
+
+    Args:
+        augments: List of (name, kwargs) tuples referencing AUGMENTS registry.
+        protected_nodes: Optional tensor of node indices that must not be removed
+            or dropped. Passed to every augmentation function; functions that do
+            not use it accept and ignore the parameter.
+    """
+    for name, kwargs in augments:
+        data = AUGMENTS.build(name, data=data, protected_nodes=protected_nodes, **kwargs)
+    return data
+
+
+class MultiView:
+    """Generate n independent augmented views of a Data object.
+
+    Supports both registry-style augments (list of (name, kwargs) tuples) and
+    class-based transforms from augmentation/transforms.py (callables Data → Data).
+
+    Args:
+        transforms: Either a list of (name, kwargs) tuples OR a list of callables.
+        n_views: Number of independent views to generate.
+    """
+
+    def __init__(self, transforms: list, n_views: int = 2):
+        self.transforms = transforms
+        self.n_views = n_views
+        # Detect which API is being used
+        self._registry_style = (
+            transforms and isinstance(transforms[0], tuple) and isinstance(transforms[0][0], str)
+        )
+
+    def __call__(
+        self,
+        data: Data,
+        protected_nodes: Optional[torch.Tensor] = None,
+    ) -> List[Data]:
+        if self._registry_style:
+            return [
+                compose(data, self.transforms, protected_nodes=protected_nodes)
+                for _ in range(self.n_views)
+            ]
+        views = []
+        for _ in range(self.n_views):
+            view = data
+            for t in self.transforms:
+                view = t(view)
+            views.append(view)
+        return views

graphssl/augmentation/functional.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+from torch_geometric.data import Data
+from torch_geometric.utils import k_hop_subgraph
+
+from graphssl.registry import AUGMENTS
+
+
+@AUGMENTS.register("edge_drop")
+def edge_drop(data: Data, p=0.2, protected_nodes: Optional[Tensor] = None):
+    assert data.edge_index is not None
+    mask = torch.rand(data.edge_index.size(1), device=data.edge_index.device) > p
+    out = data.clone()
+    out.edge_index = data.edge_index[:, mask]
+    if data.edge_attr is not None:
+        out.edge_attr = data.edge_attr[mask]
+    return out
+
+
+@AUGMENTS.register("edge_add")
+def edge_add(data: Data, p=0.1, protected_nodes: Optional[Tensor] = None):
+    assert data.edge_index is not None
+    assert data.num_nodes is not None
+    n = data.num_nodes
+    n_add = max(1, int(data.edge_index.size(1) * p))
+    src = torch.randint(0, n, (n_add,), device=data.edge_index.device)
+    dst = torch.randint(0, n, (n_add,), device=data.edge_index.device)
+    out = data.clone()
+    out.edge_index = torch.cat([data.edge_index, torch.stack([src, dst])], dim=1)
+    if data.edge_attr is not None:
+        # Synthetic edges have no known features — pad with zeros.
+        pad = torch.zeros(
+            n_add, data.edge_attr.size(1), dtype=data.edge_attr.dtype, device=data.edge_attr.device
+        )
+        out.edge_attr = torch.cat([data.edge_attr, pad], dim=0)
+    return out
+
+
+@AUGMENTS.register("subgraph")
+def subgraph(data: Data, num_hops=2, protected_nodes: Optional[Tensor] = None):
+    assert data.edge_index is not None
+    assert data.num_nodes is not None
+    n = data.num_nodes
+    seed = int(torch.randint(0, n, (1,), device=data.edge_index.device).item())
+    node_idx, edge_index_sub, _, edge_mask = k_hop_subgraph(
+        node_idx=seed,
+        num_hops=num_hops,
+        edge_index=data.edge_index,
+        num_nodes=n,
+        relabel_nodes=True,
+    )
+    out = data.clone()
+    out.x = data.x[node_idx] if data.x is not None else None
+    out.edge_index = edge_index_sub
+    out.batch = data.batch[node_idx] if data.batch is not None else None
+    if data.edge_attr is not None:
+        out.edge_attr = data.edge_attr[edge_mask]
+    return out
+
+
+@AUGMENTS.register("feat_mask")
+def feat_mask(data: Data, p=0.2, protected_nodes: Optional[Tensor] = None):
+    assert data.x is not None
+    mask = (torch.rand(data.x.size(1), device=data.x.device) > p).float()
+    out = data.clone()
+    out.x = data.x * mask
+    return out
+
+
+@AUGMENTS.register("feat_noise")
+def feat_noise(data: Data, std=0.1, protected_nodes: Optional[Tensor] = None):
+    assert data.x is not None
+    out = data.clone()
+    out.x = data.x + torch.randn_like(data.x) * std
+    return out
+
+
+@AUGMENTS.register("feat_shuffle")
+def feat_shuffle(data: Data, p=0.1, protected_nodes: Optional[Tensor] = None):
+    assert data.x is not None
+    n = data.x.size(0)
+    perm = torch.randperm(n, device=data.x.device)
+    node_mask = torch.rand(n, device=data.x.device) < p
+    x = data.x.clone()
+    x[node_mask] = data.x[perm[node_mask]]
+    out = data.clone()
+    out.x = x
+    return out
+
+
+@AUGMENTS.register("node_drop")
+def node_drop(data: Data, p=0.1, protected_nodes: Optional[Tensor] = None):
+    """Drop nodes with probability p, remapping edge_index.
+
+    Nodes in protected_nodes are never dropped (used to preserve mini-batch
+    seed nodes in NeighborLoader training).
+    """
+    assert data.x is not None
+    assert data.edge_index is not None
+    n = data.num_nodes
+    assert n is not None
+
+    keep_mask = torch.rand(n, device=data.x.device) > p
+    if protected_nodes is not None:
+        keep_mask[protected_nodes] = True
+
+    keep_idx = keep_mask.nonzero(as_tuple=True)[0]
+
+    # Build a remapping table: old node id → new node id (-1 = dropped)
+    new_idx = torch.full((n,), -1, dtype=torch.long, device=data.x.device)
+    new_idx[keep_idx] = torch.arange(keep_idx.size(0), device=data.x.device)
+
+    # Keep only edges where both endpoints survive
+    src, dst = data.edge_index
+    edge_mask = keep_mask[src] & keep_mask[dst]
+
+    out = data.clone()
+    out.x = data.x[keep_idx]
+    out.edge_index = new_idx[data.edge_index[:, edge_mask]]
+    out.batch = data.batch[keep_idx] if data.batch is not None else None
+    if data.edge_attr is not None:
+        out.edge_attr = data.edge_attr[edge_mask]
+    return out

graphssl/augmentation/transforms.py

@@ -0,0 +1,84 @@
+"""Class-based graph augmentation transforms (torchvision style).
+
+Each transform is a callable Data → Data that wraps the functional API.
+They satisfy the BaseAugmentation Protocol and can be used with Compose/MultiView.
+"""
+
+from torch_geometric.data import Data
+
+from graphssl.augmentation import functional as F
+
+
+class EdgeDrop:
+    """Randomly drop edges with probability p."""
+
+    def __init__(self, p: float = 0.2):
+        self.p = p
+
+    def __call__(self, data: Data) -> Data:
+        return F.edge_drop(data, p=self.p)
+
+
+class EdgeAdd:
+    """Randomly add edges (fraction p of existing edges)."""
+
+    def __init__(self, p: float = 0.1):
+        self.p = p
+
+    def __call__(self, data: Data) -> Data:
+        return F.edge_add(data, p=self.p)
+
+
+class Subgraph:
+    """Keep a k-hop subgraph around a random seed node."""
+
+    def __init__(self, num_hops: int = 2):
+        self.num_hops = num_hops
+
+    def __call__(self, data: Data) -> Data:
+        return F.subgraph(data, num_hops=self.num_hops)
+
+
+class FeatMask:
+    """Zero-out a fraction p of feature dimensions."""
+
+    def __init__(self, p: float = 0.2):
+        self.p = p
+
+    def __call__(self, data: Data) -> Data:
+        return F.feat_mask(data, p=self.p)
+
+
+class FeatNoise:
+    """Add Gaussian noise with standard deviation std to node features."""
+
+    def __init__(self, std: float = 0.1):
+        self.std = std
+
+    def __call__(self, data: Data) -> Data:
+        return F.feat_noise(data, std=self.std)
+
+
+class FeatShuffle:
+    """Randomly swap features between a fraction p of nodes."""
+
+    def __init__(self, p: float = 0.1):
+        self.p = p
+
+    def __call__(self, data: Data) -> Data:
+        return F.feat_shuffle(data, p=self.p)
+
+
+class NodeDrop:
+    """Drop nodes with probability p, remapping edge_index.
+
+    Note: when used with protected_nodes (seed nodes in mini-batch training),
+    use the registry-style compose() API which propagates protected_nodes
+    automatically. The class-based API does not support protected_nodes.
+    """
+
+    def __init__(self, p: float = 0.1):
+        self.p = p
+
+    def __call__(self, data: Data) -> Data:
+        return F.node_drop(data, p=self.p)

graphssl/config/__init__.py

@@ -0,0 +1,14 @@
+from .load import LOADERS, build_model, load_config
+from .schema import (
+    AFGRLConfig,
+    AugmentConfig,
+    BarlowTwinsConfig,
+    BGRLConfig,
+    DGIConfig,
+    EncoderConfig,
+    GraphCLConfig,
+    GraphDINOConfig,
+    HeadConfig,
+    SupervisedConfig,
+    VICRegConfig,
+)

graphssl/config/load.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional, Union
+
+import torch.nn as nn
+from torch_geometric.loader import DataLoader, NeighborLoader
+
+from graphssl.registry import LOADERS
+
+
+@LOADERS.register("graph")
+def build_graph_loader(*, dataset, params):
+    # Standard graph-level loader, batches whole graphs together.
+    return DataLoader(
+        dataset,
+        batch_size=params.get("batch_size", 32),
+        shuffle=params.get("shuffle", True),
+        num_workers=params.get("num_workers", 0),
+    )
+
+
+@LOADERS.register("neighbor")
+def build_neighbor_loader(*, dataset, params):
+    # Samples neighbor subgraphs per node; useful for large single-graph datasets.
+    return NeighborLoader(
+        dataset,
+        num_neighbors=params.get("num_neighbors", [10, 10]),
+        batch_size=params.get("batch_size", 1024),
+        shuffle=True,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Config loading and model construction
+# ---------------------------------------------------------------------------
+
+
+def load_config(path: Union[str, Path]) -> dict:
+    """Load a GraphSSL YAML config file and return it as a plain dict.
+
+    The returned dict can be passed directly to ``build_model()`` or to any
+    model constructor (``ModelClass(config, in_channels)``).
+
+    Requires PyYAML: ``pip install pyyaml``
+    """
+    try:
+        import yaml
+    except ImportError as e:
+        raise ImportError("PyYAML is required to load config files: pip install pyyaml") from e
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def build_model(
+    config: dict,
+    in_channels: int,
+    num_classes: Optional[int] = None,
+) -> nn.Module:
+    """Instantiate a GraphSSL model from a config dict.
+
+    ``config`` must contain a top-level ``name`` key identifying the model.
+    For ``supervised``, ``num_classes`` is required. All other fields are
+    forwarded to the model constructor and validated by its config dataclass.
+
+    Example::
+
+        cfg = load_config("configs/bgrl.yaml")
+        model = build_model(cfg, in_channels=dataset.num_features)
+
+        cfg = load_config("configs/supervised.yaml")
+        model = build_model(cfg, in_channels=dataset.num_features,
+                            num_classes=dataset.num_classes)
+
+    Args:
+        config: Config dict with a ``name`` key and model-specific fields.
+        in_channels: Node feature dimensionality from the dataset.
+        num_classes: Required only for the ``supervised`` model.
+
+    Returns:
+        Instantiated model as an ``nn.Module``.
+    """
+    # Deferred import to avoid circular dependency (models import from config).
+    from graphssl.models import (
+        AFGRL,
+        BGRL,
+        DGI,
+        BarlowTwins,
+        GraphCL,
+        GraphDINO,
+        Supervised,
+        VICReg,
+    )
+
+    _MODELS = {
+        "dgi": DGI,
+        "graphcl": GraphCL,
+        "vicreg": VICReg,
+        "barlow_twins": BarlowTwins,
+        "bgrl": BGRL,
+        "afgrl": AFGRL,
+        "graphdino": GraphDINO,
+        "supervised": Supervised,
+    }
+
+    name = config.get("name", "").lower().replace("-", "_")
+    if name not in _MODELS:
+        raise ValueError(f"Unknown model '{name}'. Available: {sorted(_MODELS)}")
+    cls = _MODELS[name]
+    if name == "supervised":
+        if num_classes is None:
+            raise ValueError("num_classes is required for the Supervised model")
+        return cls(config, in_channels, num_classes)
+    return cls(config, in_channels)