Configure and Train a GNN

The GNN class accepts a number of optional hyperparameters that fine-tune training, the network architecture, and other behaviors. Each parameter has a sensible default, so you only need to set the ones relevant to your task. The sections below group them by role.

Set general parameters

Parameter	Type	Description
`device`	`str`	Device to perform training, inference, and feature extraction. One of `"cuda"` or `"cpu"`. Default: `"cuda"`. If your predictive reasoner instance does not have a GPU, the system automatically falls back to `"cpu"`, so this setting is safe to leave at its default.
`seed`	`int`	Random seed for reproducibility. Default: `42`.

Set training parameters

Parameter	Type	Description
`n_epochs`	`int`	Number of training epochs. An epoch corresponds to a full pass over the training data. Default: `10`
`max_iters`	`int`	Maximum number of batch iterations per epoch. If `None`, all batches are processed. Otherwise, limits iterations. Default: `None`.
`train_batch_size`	`int`	Batch size for training. Default: `128`.
`val_batch_size`	`int`	Batch size for validation. Default: `128`.
`eval_every`	`int`	Frequency (in epochs) to evaluate on the validation set. Default: `1`.
`patience`	`int`	Number of epochs without improvement before early stopping. Default: `5`.
`lr`	`float`	Learning rate. Default: `0.001`.
`T_max`	`int`	Max iterations for cosine annealing scheduler. Defaults to `n_epochs` if `None`. Default: `None`.
`eta_min`	`int`	Minimum learning rate for cosine annealing. Default: `1e-8`.

Configure labels and loss

Parameter	Type	Description
`label_smoothing`	`bool`	Whether to apply label smoothing (for classification). Default: `False`.
`label_smoothing_alpha`	`float`	Smoothing parameter α ∈ (0, 1). Default: `0.1`.
`clamp_min`	`int`	Specifies the lower bound of the model’s output distribution in percentile terms (0–100). A value of 0 means no lower percentile cutoff is applied, while higher values restrict predictions to exclude the lowest portion of the output distribution. Default: `0`.
`clamp_max`	`int`	Specifies the higher bound of the model’s output distribution in percentile terms (0–100). A value of 100 means no higher percentile cutoff is applied, while lower values restrict predictions to exclude the highest portion of the output distribution. Default: `100`.

Configure the GNN architecture

Parameter	Type	Description
`channels`	`int`	Hidden channels for GNN, encoders, and prediction heads. Default: `128`.
`gnn_layers`	`int`	Number of GNN layers. Defaults to `len(fanouts)` if `None`. Default: `None`.
`fanouts`	`List[int]`	Neighbors to sample per GNN layer. E.g., `[128, 64]`. Default: `[128, 64]`.
`conv_aggregation`	`str`	Aggregation method for convolutions. It can be one of `"mean"`, `"max"` or `"sum"`. Default: `"mean"`.
`hetero_conv_aggregation`	`str`	Aggregation across edge types in heterogeneous graphs. It can be one of `"mean"`, `"max"` or `"sum"`. Default: `"sum"`.
`gnn_norm`	`str`	Normalization for GNN layers. It can be one of `"batch_norm"`, `"layer_norm"` or `"instance_norm"`. Default: `"layer_norm"`.

Configure the prediction head

Parameter	Type	Description
`head_layers`	`int`	Number of Multi-Layer Perceptron (MLP) layers in the prediction head. Default: `1`.
`head_norm`	`str`	Normalization for the MLP prediction head. It can be one of `"batch_norm"` or `"layer_norm"`. Default: `"batch_norm"`.

Configure temporal sampling

Parameter	Type	Description
`use_temporal_encoder`	`bool`	Whether to use a temporal encoding model. Default: `True`.
`temporal_strategy`	`str`	Strategy for temporal neighbor sampling. `"uniform"` ignores time; `"last"` picks most recent. Default: `"uniform"`.

Configure negative sampling for link prediction

Parameter	Type	Description
`num_negative`	`int`	Number of negative samples per source node (for link prediction). Default: `10`.
`negative_sampling_strategy`	`str`	Strategy: `"random"` or `"degree_based"`. `"degree_based"` favors popular nodes. Default: `"random"`.

Configure embeddings and shallow features

Parameter	Type	Description
`text_embedder`	`str`	Text embedding model. It can be one of `"model2vec-potion-base-4M"` or `"bert-base-distill"`. Default: `"model2vec-potion-base-4M"`.
`id_awareness`	`bool`	Whether to use ID-awareness embeddings. Default: `False`.
`shallow_embeddings_list`	`List[str]`	Tables to assign learnable shallow embeddings. Default: `[]`.

Example

Hyperparameters are passed as keyword arguments when constructing the GNN. A common pattern is to collect them in a dictionary and unpack it with **, which keeps the tuning knobs separate from the structural arguments and makes the configuration easy to tweak between runs:

train_config = {
    "device": "cuda",
    "n_epochs": 10,
    "train_batch_size": 256,
    "lr": 0.001,
    "head_layers": 2,
    "label_smoothing": True,
    "patience": 5,
}

gnn = GNN(
    exp_database="EXPERIMENTS_DB",
    exp_schema="MODEL_REGISTRY",
    graph=gnn_graph,
    property_transformer=pt,
    train=Train,
    validation=Validation,
    task_type="binary_classification",
    eval_metric="roc_auc",
    **train_config,
)
gnn.fit()

You can also pass each hyperparameter directly as a keyword argument — the dictionary is just a convenience.

Train the GNN

Once the GNN is configured, call .fit() to start training:

gnn.fit()