Config API

EncoderConfig dataclass

EncoderConfig(name: str, hidden_dim: int, num_layers: int, mlp_ratio: float = 2.0, drop: Optional[float] = None, pool: bool = True, readout: str = 'mean', norm_type: str = 'batch', weight_standardization: bool = False, edge_dim: Optional[int] = None, node_emb_num_classes: Optional[int] = None, edge_emb_num_classes: Optional[int] = None)

Validated configuration for graph encoders.

All three encoder backbones (GCN, GIN, Transformer) share this config: build() forwards to an encoder the fields its constructor takes. A field left at its default is simply not forwarded to an encoder that has no use for it; a field that was set and that the chosen encoder does not take makes build() raise (see there).

build

build(in_channels: int) -> nn.Module

Instantiate the encoder via the ENCODERS registry.

Only the kwargs accepted by the encoder's init are forwarded, so the config can carry fields that an encoder has no use for (mlp_ratio with GCN, …). A field that was set — a required one, or one whose value differs from its default — and that the encoder does not take raises ValueError: it would otherwise be discarded without notice, as num_layers once was with the GCN encoder, drop with the GCN and Transformer encoders, and as edge_dim would be with an encoder that cannot use edge features. pool and readout are exempt: the models read them, not the encoder. An encoder whose constructor has **kwargs receives every field.

Source code in src/graphssl/config/schema.py

def build(self, in_channels: int) -> nn.Module:
    """Instantiate the encoder via the ENCODERS registry.

    Only the kwargs accepted by the encoder's __init__ are forwarded, so the config
    can carry fields that an encoder has no use for (``mlp_ratio`` with GCN, …).
    A field that was set — a required one, or one whose value differs from its
    default — and that the encoder does not take raises ``ValueError``: it would
    otherwise be discarded without notice, as ``num_layers`` once was with the GCN
    encoder, ``drop`` with the GCN and Transformer encoders, and as ``edge_dim``
    would be with an encoder that cannot use edge features.
    ``pool`` and ``readout`` are exempt: the models read them, not the encoder.
    An encoder whose constructor has ``**kwargs`` receives every field.
    """
    import inspect

    from graphssl.registry import ENCODERS

    cls = ENCODERS.get_builder(self.name)
    params = inspect.signature(cls.__init__).parameters
    takes_any = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
    valid = set(params) - {"self"}
    if not takes_any:
        ignored = [
            f.name
            for f in fields(self)
            if f.name not in valid
            and f.name not in ("name", "pool", "readout")
            # a required field has no default: it is always set
            and (f.default is MISSING or getattr(self, f.name) != f.default)
        ]
        if ignored:
            raise ValueError(
                f"encoder {self.name!r} does not take {', '.join(ignored)}: the value "
                "would be discarded. Remove it from the encoder config, choose an "
                "encoder that supports it, or make the encoder's constructor accept it"
            )
    kwargs = {k: v for k, v in vars(self).items() if k != "name" and (takes_any or k in valid)}
    if kwargs.get("drop") is None:
        kwargs.pop("drop", None)  # None: the encoder's own default
    return ENCODERS.build(self.name, in_channels=in_channels, **kwargs)

SupervisedConfig dataclass

SupervisedConfig(encoder: EncoderConfig, task: str = 'classification')

Validated configuration for the supervised baseline.

DGIConfig dataclass

DGIConfig(encoder: EncoderConfig, corruption: str = 'shuffle_nodes', shuffle_ratio: float = 1.0)

Validated configuration for DGI.

GraphCLConfig dataclass

GraphCLConfig(encoder: EncoderConfig, augment: List[AugmentConfig] = list(), proj_dim: int = 128, tau: float = 0.5, loss_chunk_size: Optional[int] = 4096)

Validated configuration for GraphCL.

BGRLConfig dataclass

BGRLConfig(encoder: EncoderConfig, augment: List[AugmentConfig] = list(), pred_hidden: int = 512, pred_norm: str = 'batch', pred_activation: str = 'relu', ema_tau: float = 0.99, ema_tau_end: float = 1.0, total_steps: int = 0)

Validated configuration for BGRL.

AFGRLConfig dataclass

AFGRLConfig(encoder: EncoderConfig, pred_hidden: int = 512, pred_norm: str = 'batch', pred_activation: str = 'relu', ema_tau: float = 0.99, ema_tau_end: float = 1.0, total_steps: int = 0, topk: int = 5, num_centroids: int = 50, num_kmeans: int = 4, clus_num_iters: int = 20, kmeans_threads: Optional[int] = 8, knn_chunk_size: Optional[int] = 4096)

Validated configuration for AFGRL.

VICRegConfig dataclass

VICRegConfig(encoder: EncoderConfig, augment: List[AugmentConfig] = list(), proj_dim: int = 256, invariance: float = 25.0, variance: float = 25.0, covariance: float = 1.0)

Validated configuration for VICReg.

BarlowTwinsConfig dataclass

BarlowTwinsConfig(encoder: EncoderConfig, augment: List[AugmentConfig] = list(), proj_dim: int = 256, lambda_param: Optional[float] = None)

Validated configuration for Barlow Twins.

GraphDINOConfig dataclass

GraphDINOConfig(encoder: EncoderConfig, head: HeadConfig, augment_teacher: List[AugmentConfig] = list(), augment_student: List[AugmentConfig] = list(), ema_tau: float = 0.996, ema_tau_base: float = 0.9, total_steps: int = 0, freeze_last_layer_epochs: int = 1, n_views: int = 2, n_global_views: int = 2)

Validated top-level configuration for GraphDINO.

HeadConfig dataclass

HeadConfig(name: str, proj_hidden: int, bottleneck_dim: int, n_prototypes: int, student_temp: float = 0.1, teacher_temp: float = 0.07, center_momentum: float = 0.9, warmup_teacher_temp: float = 0.04, warmup_teacher_temp_epochs: int = 30, norm_last_layer: bool = False)

Validated configuration for the DINOHead projection head.

load_config

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

Source code in src/graphssl/config/load.py

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)

build_model

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)

Parameters:

Name	Type	Description	Default
config	dict	Config dict with a name key and model-specific fields.	required
in_channels	int	Node feature dimensionality from the dataset.	required
num_classes	Optional[int]	Required only for the supervised model.	None

Returns:

Type	Description
Module	Instantiated model as an nn.Module.

Source code in src/graphssl/config/load.py

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)