TrainerConfig

Defines training parameters.

Parameters

General Settings

Parameter	Type	Description	Optional
`connector`	`SnowflakeConnector`	The connector object used for sending requests to the GNN engine.	No
`experiment_config`	`ExperimentConfig`	An instance of `ExperimentConfig` which defines the database and schema that experiment metadata will be saved.	No
`device`	`str`	Device to perform training, inference, and feature extraction. One of `"cuda"` or `"cpu"`.	No
`seed`	`int`	Random seed for reproducibility. Default: `42`.	Yes

Training Settings

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

Label & Loss Settings

Parameter	Type	Description	Optional
`label_smoothing`	`bool`	Whether to apply label smoothing (for classification). Default: `False`.	Yes
`label_smoothing_alpha`	`float`	Smoothing parameter α ∈ (0, 1). Default: `0.1`.	Yes
`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`.	Yes
`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`.	Yes

Graph Neural Network (GNN) Settings

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

Prediction Head Settings

Parameter	Type	Description	Optional
`head_layers`	`int`	Number of MLP layers in the prediction head. Default: `1`.	Yes
`head_norm`	`str`	Normalization for the MLP prediction head. It can be one of “batch_norm” or “layer_norm”. Default: `"batch_norm"`.	Yes

Temporal Settings

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

Negative Sampling (Link Prediction)

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

Embeddings & Shallow Features

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

Returns

An instance of the TrainerConfig class.

Example

from relationalai_gnns import ExperimentConfig, TrainerConfig

experiment_config = ExperimentConfig(database="database_name",
                                     schema="schema_name")

trainer_config = TrainerConfig(
    connector=connector,
    experiment_config=experiment_config,
    device="cuda",
    n_epochs=10,
    patience=5)

Methods

Name	Description	Returns
`to_dict`	Writes `TrainerConfig` contents in a dictionary	`Dict`
`validate`	Validates the model’s configuration	`bool`

.to_dict()

Writes TrainerConfig contents in a dictionary.

model_config_dict = config.to_dict()

.validate()

Validates the model configuration. Returns True if the configuration is valid; otherwise, returns False along with a list of validation errors.

config.validate()