Encoders API

GCNEncoder

GCNEncoder(in_channels: int, hidden_dim: int, out_dim: int | None = None, norm_type: str = 'batch', weight_standardization: bool = False, batchnorm_mm: float = 0.1, pool: bool = False, batchnorm: bool | None = None, layernorm: bool | None = None, num_layers: int = 2, drop: float = 0.0)

Bases: Module

GCNConv backbone (2 layers by default).

Architecture per layer: (Dropout →) GCNConv → Norm → PReLU.

Uses the same norm_type API as GINEncoder and TransformerEncoder so that switching encoder backbones is a one-line config change.

Parameters:

Name	Type	Description	Default
in_channels	int	Input node feature dimension.	required
hidden_dim	int	Output dimension of every layer but the last.	required
out_dim	int | None	Output dimension of the last layer (defaults to hidden_dim).	None
norm_type	str	'batch', 'layer', or 'none' — same API as GINEncoder.	'batch'
weight_standardization	bool	Standardize the GCNConv weights of every layer after the first, before each forward (BGRL's setting on ogbn-arxiv, with layer norm).	False
batchnorm_mm	float	BatchNorm momentum (standard PyTorch convention; default 0.1, as in the other encoders). With 0.01 the running statistics still hold 5% of their initial values (variance 1) after 300 forward passes, where the batch variance of a first GCN layer is 1e-4 to 1e-2: eval-mode outputs then differ from the ones the model was trained with.	0.1
pool	bool	If True, apply global_mean_pool to produce graph-level embeddings.	False
num_layers	int	Number of GCNConv layers.	2
drop	float	Dropout on the input of every layer, as in Kipf & Welling (2017). 0 by default: the encoder had no dropout before this parameter existed, and the drop of a config was discarded without notice.	0.0

Source code in src/graphssl/encoders/gcn.py

def __init__(
    self,
    in_channels: int,
    hidden_dim: int,
    out_dim: int | None = None,
    norm_type: str = "batch",
    weight_standardization: bool = False,
    batchnorm_mm: float = 0.1,
    pool: bool = False,
    # Legacy aliases kept for backwards compatibility — map to norm_type internally.
    batchnorm: bool | None = None,
    layernorm: bool | None = None,
    num_layers: int = 2,
    drop: float = 0.0,
):
    super().__init__()
    if out_dim is None:
        out_dim = hidden_dim
    if num_layers <= 0:
        raise ValueError(f"num_layers must be > 0, got {num_layers}")
    if not (0.0 <= drop < 1.0):
        raise ValueError(f"drop must be in [0, 1), got {drop}")

    # --- Backwards-compatible norm resolution ---
    # If the caller uses the old batchnorm/layernorm bool flags, honour them.
    if batchnorm is not None or layernorm is not None:
        if batchnorm and layernorm:
            raise ValueError("batchnorm and layernorm are mutually exclusive.")
        if batchnorm:
            norm_type = "batch"
        elif layernorm:
            norm_type = "layer"
        else:
            norm_type = "none"

    self.weight_standardization = weight_standardization
    self.pool = pool
    self.drop = drop

    dims = [in_channels] + [hidden_dim] * (num_layers - 1) + [out_dim]
    self.convs = nn.ModuleList(GCNConv(i, o) for i, o in zip(dims[:-1], dims[1:], strict=True))
    self.norms = nn.ModuleList(_make_norm(norm_type, d, batchnorm_mm) for d in dims[1:])
    self.acts = nn.ModuleList(nn.PReLU() for _ in dims[1:])

GINEncoder

GINEncoder(in_channels: int, hidden_dim: int, out_dim: int | None = None, num_layers: int = 3, mlp_ratio: int = 2, drop: float = 0.2, pool: bool = False, norm_type: str = 'batch', edge_dim: int | None = None, node_emb_num_classes: int | None = None, edge_emb_num_classes: int | None = None)

Bases: Module

GIN backbone with configurable norm, optional edge features, and categorical embeddings.

Supports both continuous and categorical node/edge features out of the box, making it compatible with datasets like ZINC (categorical atoms + bond types) and ogbn-arxiv (continuous 128-dim embeddings).

Parameters:

Name	Type	Description	Default
in_channels	int	Input node feature dimension. Ignored when node_emb_num_classes is set.	required
hidden_dim	int	Hidden dimension throughout the stack.	required
out_dim	int | None	Final output dimension (defaults to hidden_dim).	None
num_layers	int	Number of GINLayer blocks.	3
mlp_ratio	int	Width multiplier for the inner MLP of each GINConv.	2
drop	float	Dropout rate.	0.2
pool	bool	If True, apply global_add_pool to produce graph-level embeddings.	False
norm_type	str	Normalisation after each layer — 'batch', 'layer', or 'none'.	'batch'
edge_dim	int | None	Projected edge feature dimension fed to GINEConv. Required when edge features are present. If None, edge_attr is ignored.	None
node_emb_num_classes	int | None	If set, replaces the linear input projection with nn.Embedding (needed for integer node labels such as ZINC atom types, 28 classes).	None
edge_emb_num_classes	int | None	If set, adds an nn.Embedding that maps integer edge labels to edge_dim-dimensional vectors before GINEConv. Required for ZINC bond types (4 classes). Needs edge_dim to be set.	None

Source code in src/graphssl/encoders/gin.py

def __init__(
    self,
    in_channels: int,
    hidden_dim: int,
    out_dim: int | None = None,
    num_layers: int = 3,
    mlp_ratio: int = 2,
    drop: float = 0.2,
    pool: bool = False,
    norm_type: str = "batch",
    edge_dim: int | None = None,
    node_emb_num_classes: int | None = None,
    edge_emb_num_classes: int | None = None,
):
    super().__init__()
    if out_dim is None:
        out_dim = hidden_dim
    if edge_emb_num_classes is not None and edge_dim is None:
        raise ValueError("edge_emb_num_classes requires edge_dim to be set.")
    self.pool = pool
    self._drop = drop
    self.edge_dim = edge_dim

    # --- Node input projection ---
    # Categorical (e.g. ZINC atom types)  → nn.Embedding
    # Continuous (e.g. ogbn-arxiv 128-d)  → nn.Linear
    if node_emb_num_classes is not None:
        self.input_proj: nn.Module = nn.Embedding(node_emb_num_classes, hidden_dim)
    else:
        self.input_proj = Linear(in_channels, hidden_dim)

    # --- Edge input projection ---
    # Categorical (e.g. ZINC bond types 0-3) → nn.Embedding → edge_dim vectors
    # Continuous edge features               → passed directly (user must ensure dim matches)
    # No edge features                        → None
    self.edge_proj: nn.Module | None = None
    if edge_emb_num_classes is not None:
        assert edge_dim is not None  # validated above: edge_emb_num_classes requires edge_dim
        self.edge_proj = nn.Embedding(edge_emb_num_classes, edge_dim)

    self.layers = nn.ModuleList(
        [
            GINLayer(
                hidden_dim,
                mlp_ratio=mlp_ratio,
                drop=drop,
                norm_type=norm_type,
                edge_dim=edge_dim,
            )
            for _ in range(num_layers)
        ]
    )
    self.lin1 = Linear(hidden_dim, hidden_dim)
    self.lin2 = Linear(hidden_dim, out_dim)
    self.activation = ReLU()

TransformerEncoder

TransformerEncoder(in_channels: int, hidden_dim: int, num_layers: int = 4, heads: int = 4, dropout: float = 0.1, attn_dropout: float = 0.1, mlp_ratio: float = 4.0, concat: bool = True, norm_type: str = 'batch', act_layer_name: str = 'gelu', pool: bool = True, edge_dim: int | None = None, node_emb_num_classes: int | None = None, edge_emb_num_classes: int | None = None, drop: float | None = None)

Bases: Module

Stack of pre-norm TransformerConv blocks.

Aligns with the GINEncoder API: same norm_type values ('batch', 'layer', 'none'), same edge_dim / node_emb_num_classes / edge_emb_num_classes parameters for seamless config swap between encoder types.

Parameters:

Name	Type	Description	Default
in_channels	int	Input node feature size. Ignored when node_emb_num_classes is set.	required
hidden_dim	int	Uniform hidden/output dimension throughout the stack.	required
num_layers	int	Number of TransformerBlock layers.	4
heads	int	Number of attention heads (hidden_dim must be divisible by heads).	4
dropout	float	Dropout in the FFN.	0.1
attn_dropout	float	Dropout inside TransformerConv attention.	0.1
mlp_ratio	float	FFN hidden dim = hidden_dim * mlp_ratio.	4.0
concat	bool	If True, TransformerConv concatenates heads (see TransformerBlock).	True
norm_type	str	'batch', 'layer', or 'none' — aligned with GINEncoder API.	'batch'
act_layer_name	str	'gelu' or 'relu'.	'gelu'
pool	bool	Apply global_mean_pool to produce graph-level embeddings.	True
edge_dim	int | None	Projected edge feature dimension. When set, each TransformerBlock conditions attention on edge_attr. Ignored if None.	None
node_emb_num_classes	int | None	If set, replaces the linear input projection with nn.Embedding (for integer node labels, e.g. ZINC atom types).	None
edge_emb_num_classes	int | None	If set, adds an nn.Embedding that maps integer edge labels to edge_dim vectors before attention. Requires edge_dim to be set.	None
drop	float | None	The single dropout rate of EncoderConfig: when given, it replaces both dropout and attn_dropout (drop=0 means no dropout anywhere).	None

Source code in src/graphssl/encoders/transformer.py

def __init__(
    self,
    in_channels: int,
    hidden_dim: int,
    num_layers: int = 4,
    heads: int = 4,
    dropout: float = 0.1,
    attn_dropout: float = 0.1,
    mlp_ratio: float = 4.0,
    concat: bool = True,
    norm_type: str = "batch",
    act_layer_name: str = "gelu",
    pool: bool = True,
    edge_dim: int | None = None,
    node_emb_num_classes: int | None = None,
    edge_emb_num_classes: int | None = None,
    drop: float | None = None,
):
    super().__init__()
    assert hidden_dim % heads == 0, "hidden_dim must be divisible by heads"
    if drop is not None:
        if not (0.0 <= drop < 1.0):
            raise ValueError(f"drop must be in [0, 1), got {drop}")
        dropout = attn_dropout = drop
    if edge_emb_num_classes is not None and edge_dim is None:
        raise ValueError("edge_emb_num_classes requires edge_dim to be set.")

    self.pool = pool
    self.edge_dim = edge_dim

    # --- Node input projection ---
    if node_emb_num_classes is not None:
        self.input_proj: nn.Module = nn.Embedding(node_emb_num_classes, hidden_dim)
    else:
        self.input_proj = nn.Linear(in_channels, hidden_dim)

    # --- Edge input projection ---
    self.edge_proj: nn.Module | None = None
    if edge_emb_num_classes is not None:
        assert edge_dim is not None  # validated above: edge_emb_num_classes requires edge_dim
        self.edge_proj = nn.Embedding(edge_emb_num_classes, edge_dim)

    self.blocks = nn.ModuleList(
        [
            TransformerBlock(
                dim=hidden_dim,
                heads=heads,
                dropout=dropout,
                attn_dropout=attn_dropout,
                mlp_ratio=mlp_ratio,
                concat=concat,
                norm_type=norm_type,
                act_layer_name=act_layer_name,
                edge_dim=edge_dim,
            )
            for _ in range(num_layers)
        ]
    )
    self.norm = _make_norm(norm_type, hidden_dim)

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)