dgf.learning

dgf.learning.LinkPredictionModel

dgf.learning.Model

Bases: ABC

A generic model from the 10-lines of code API.

A model is a high-level, user-facing object that "makes predictions". Practically, the model can encapsulate core GNN models, sampler configurations, normalization settings, padding configurations, and any other data required to run the model on raw user data.

Each Model subclass must provide a unique identifier via its name() class method. This identifier is used for registration, saving, and restoring models.

A model can be saved on disk with the "model.save(path)" method. A saved and reloaded model is exactly equivalent to the original model (no lost information; this is not an export). The model data is composed of 3 artifacts: - The "metadata.json" file that contains generic information about this Model class. - The "data.json" file that contains lightweight, JSON-serializable data returned by the "model.data()" method. - The content written or read by the abstract _internal_save and _internal_load methods. This can be any data, and it is generally suited for writing large data chunks, e.g., neural network model weights.

`init`

Initializes the Model.

Subclasses should implement their specific initialization logic using the provided data.

Parameters:

Name	Type	Description	Default
`data`	`Any`	A dataclass instance containing JSON-serializable data required to construct the model. This is typically the output of the `data()` method.	required

`data` `abstractmethod`

Returns a JSON-serializable dataclass instance representing the model's data.

This data is used as the data argument in the model's constructor when loading the model from disk.

`describe` `abstractmethod`

Text or colab augmented text describing the model.

`name` `abstractmethod` `classmethod`

The unique name of the model used for serialization and registration.

`save`

Saves a model to disk. Can be later reloaded with "load_model".

Usage example

model = dgf.learning.train_node_model(...)
model.save("/tmp/my_model")
loaded_model = dgf.learning.load_model("/tmp/my_model")

Parameters:

Name	Type	Description	Default
`path`	`str`	The directory path where the model should be saved.	required

dgf.learning.NodePredictionModel

Bases: Model

The user-visible returned model object.

`evaluate`

Evaluates the model on a given graph.

Usage example:

model = train_node_model(train_graph, ...)
evaluation = model.evaluate(test_graph)
# Show the evaluation in a colab
evaluation
# Access the evaluation values
print(evaluation.accuracy)

Parameters:

Name	Type	Description	Default
`graph`	`InMemoryGraph`	The input graph data. This graph should contain the true labels for the target column.	required
`num_eval_steps`	`Optional[int]`	Maximum number of evaluation batches to run. If None, evaluation is run on all specified `seed_node_idxs`.	`10000`
`seed_node_idxs`	`Optional[SeedNodeIdxs]`	Indices of the seed nodes within the target nodeset to evaluate. If None, all nodes in the target nodeset are used.	`None`
`verbose`	`int`	The verbosity level. Higher values provide more detailed logging output during the evaluation process.	`2`
`random_seed`	`Optional[int]`	Random seed to select the seed nodes is num_train_steps is less than the number of nodes in the graph. If None, use a new random seed each time.	`None`

Returns:

Type	Description
`Evaluation`	An `evaluation.Evaluation` object containing the evaluation results,
`Evaluation`	specifically the accuracy and the number of examples evaluated.

`label_classes`

Returns the string representation of the labels.

`num_label_classes`

Returns the number of classes in the target label column.

`predict`

Predicts the target column values for the given seed nodes.

Parameters:

Name	Type	Description	Default
`graph`	`InMemoryGraph`	The input graph.	required
`seed_node_idxs`	`SeedNodeIdxs`	The indices of the seed nodes in the target nodeset.	required
`verbose`	`int`	The verbosity level.	`2`

Returns:

Type	Description
`ndarray`	An array of probabilities for each seed node.

`predict_batch`

Generate batches of predictions.

`predict_on_graph_sample_batch`

Predicts the target column values for a batch of graph samples.

Parameters:

Name	Type	Description	Default
`graph_samples`	`List[InMemoryGraph]`	A list of graph samples to predict on.	required

Returns:

Type	Description
`ndarray`	An array of probabilities for each graph sample.

`to_tensorflow_function`

Exports the model as a TensorFlow function without the sampling step.

This method creates a TensorFlow function that encapsulates the model's normalization and core prediction logic. It is designed to be used with pre-sampled subgraphs.

The returned TensorFlow function takes two arguments

graph: Either a dgf.data.TFInMemoryGraph or dgf.data.TFInMemoryGraphDict graph (depending on consume_tf_graph_dict).
seed_node_idxs: A 1D tensor of integers specifying the indices of the nodes within the target nodeset for which to generate predictions.

Parameters:

Name	Type	Description	Default
`consume_tf_graph_dict`	`bool`	If `True`, the returned TensorFlow function will expect a `dgf.data.TFInMemoryGraphDict` as the `graph` argument. This format is a flat dictionary, which is often easier to use with TF SavedModel signatures. If `False` (default), the function will expect a `dgf.data.TFInMemoryGraph` object. While more natural, this can lead to more complex manual creation of TF SavedModel signatures.	`False`

Returns:

Type	Description
`Module`	A `tf.Module` with a `__call__` method that takes `graph` and
`Module`	`seed_node_idxs` and returns a tensor of predictions for the seed nodes.

dgf.learning.load_model

Loads a model previously saved with model.save().

Usage example:

  model = dgf.learning.train_node_model(...)
  model.save("/tmp/my_model")
  loaded_model = dgf.learning.load_model("/tmp/my_model")

Parameters:

Name	Type	Description	Default
`path`	`str`	The directory path where the model was saved.	required

Returns:

Type	Description
`Model`	The loaded Model instance.

dgf.learning.train_link_model

Trains a supervised Graph Neural Network model for edge prediction.

Usage example:

model = dgf.learning.train_link_model(
    graph=graph,
    schema=schema,
    target_edgeset="cites",
)

The model is trained to predict the existence probability of an edge.

Parameters:

Name	Type	Description	Default
`graph`	`InMemoryGraph`	The input graph data structure. Training training data, as well as the validation data if `valid_graph` is not specified.	required
`schema`	`GraphSchema`	The schema of the graph.	required
`target_edgeset`	`Optional[str]`	The name of the edgeset to predict. If not specified, inferred if only one exists. The model is trained as a recommender system where the target edgset contains edges from queries to documents.	`None`
`max_training_time_seconds`	`Optional[int]`	Maximum training time in seconds. If not provided, does not limit training time.	`None`
`work_dir`	`Optional[str]`	Working directory to store checkpoints.	`None`
`verbose`	`int`	Verbosity level. From 0 (no logs) to 2 (lots of logs).	`2`
`validation_ratio`	`float`	Ratio of edges to use for validation. Only used if the "valid_graph" and "valid_seed_edges" arguments are not provided.	`0.1`
`train_seed_edges`	`Optional[List[int]]`	List of edge indices to use for training. If not provided, uses all the edges (modulo the extraction of validation edges).	`None`
`valid_seed_edges`	`Optional[List[int]]`	List of edge indices to use for validation.	`None`
`num_train_steps`	`Optional[int]`	Number of training steps.	`10000`
`num_valid_steps`	`Optional[int]`	Optional. Maximum number of validation steps. If validation caching is enabled (cache_valid_dataset=True), the same validation batch will be used each time. Otherwise, the validation batch will be sampled without replacement for each validation.	`1000`
`valid_every_n_steps`	`int`	Validate every n steps.	`1000`
`valid_graph`	`Optional[InMemoryGraph]`	An optional graph for validation. If None, `graph` is used.	`None`
`num_sampling_hops`	`int`	Number of sampling hops.	`2`
`sampling_width`	`int`	Sampling width.	`15`
`num_layers`	`int`	Number of GNN layers.	`2`
`batch_size`	`int`	Batch size.	`8`
`node_embedding_dim`	`int`	Node embedding dimension.	`128`
`learning_rate`	`float`	Learning rate.	`0.001`
`cache_valid_dataset`	`bool`	If True, the validation dataset is cached in memory.	`True`
`num_negative_nodes`	`int`	Number of negative target nodes to sample for each edge.	`8`
`message_passing_on_target_edgeset`	`bool`	If True, message passing is allowed to use edges from the `target_edgeset`. Otherwise, these edges are excluded from the GNN message passing. Note: If `message_passing_on_target_edgeset=true` the seed edge used to generate positive samples is still masked.	`True`
`negative_edges`	`Literal['random', 'random-walk']`	The strategy to use for sampling negative edges. Can be "random" or "random-walk".	`'random'`
`random_walk_num_walks_per_negative`	`int`	Number of random walks to perform per negative sample when `negative_edges` is "random-walk".	`10`
`diagnostic_dir`	`Optional[str]`	If provided, creates this directory and export to it artefacts that can be useful to understand and debug the model training.	`None`
`experimental_preprocess_core_model_config`	`Optional[Callable[[CoreModelConfig], CoreModelConfig]]`	An optional function to preprocess the `CoreModelConfig` before it is used to create the core model.	`None`
`cache_normalized_features`	`bool`	If True, pre-compute the normalized features during the preparation stage instead of computing them on the fly in the generator. This option can speed up data generation/training but increases memory consumption.	`False`
`cache_normalized_features_device`	`Literal['host', 'device']`	Specifies the device ("host" for RAM or "device" for GPU/TPU) to store the cached normalized features. Caching features on the same device used for training reduces host-device communication, potentially speeding up training, but increases memory consumption on the device.	`'device'`
`export_metrics_to_xm`	`bool`	If True, export training and validation metrics to XManager.	`False`
`architecture`	`Union[Architecture, str]`	The architecture of the GNN model to use.	`DEFAULT_ARCHITECTURE`
`source_sampling_plan`	`Optional[SamplingPlan]`	An advanced option to provide a custom plan for the sampler of the source node. When you use this option, the sampler ignores standard graph sampling arguments and validation checks e.g., num_sampling_hops, sampling_width.	`None`
`target_sampling_plan`	`Optional[SamplingPlan]`	An advanced option to provide a custom plan for the sampler of the target node. When you use this option, the sampler ignores standard graph sampling arguments and validation checks e.g., num_sampling_hops, sampling_width.	`None`

Returns:

Type	Description
`LinkPredictionModel`	A LinkPredictionModel instance.

dgf.learning.train_node_model

Trains a supervised Graph Neural Network model for node-level prediction.

This function trains a GNN to predict a specific target column within a designated nodeset of the provided graph.

Parameters:

Name	Type	Description	Default
`graph`	`Graph`	The input graph data structure.	required
`schema`	`GraphSchema`	The schema of the graph.	required
`target_column`	`str`	The name of the node feature column to be predicted.	required
`target_nodeset`	`Optional[str]`	The name of the nodeset containing the target nodes. If not provided, it's inferred if there is only one nodeset.	`None`
`max_training_time_seconds`	`Optional[int]`	Optional. The maximum duration for training in seconds. If None, training runs until convergence or default limits.	`None`
`work_dir`	`Optional[str]`	Optional. Directory to store model checkpoints. If not provided, the training is not checkpointed.	`None`
`verbose`	`int`	The verbosity level. Higher values provide more output.	`2`
`validation_ratio`	`float`	Ratio of the training dataset used to create the validation dataset in case no validation dataset is manually provided e.g., train_seed_nodes and valid_seed_nodes are provided. If set to 0, the entire dataset is used for training, and the tree is not pruned.	`0.1`
`train_seed_nodes`	`Optional[SeedNodeIdxs]`	Optional. A np.ndarray or list of integer indices specifying the subset of nodes within the `target_nodeset` to be used for training. If None, the training nodes are determined based on `validation_ratio` and `valid_nodes`.	`None`
`valid_seed_nodes`	`Optional[SeedNodeIdxs]`	Optional. A np.ndarray or list of integer indices specifying the subset of nodes within the `target_nodeset` to be used for validation. If None, the validation set is determined based on `validation_ratio` and `train_nodes`. If both `train_nodes` and `valid_nodes` are None, the data is split according to `validation_ratio`.	`None`
`num_train_steps`	`Optional[int]`	Optional. The number of training steps to perform.	`10000`
`num_valid_steps`	`Optional[int]`	Optional. Maximum number of validation steps. If validation caching is enabled (cache_valid_dataset=True), the same validation batch will be used each time. Otherwise, the validation batch will be sampled without replacement for each validation.	`1000`
`valid_every_n_steps`	`int`	The number of training steps between each validation.	`1000`
`graph_format`	`Union[GraphFormat, str]`	Optional. The format of the input graph. If set to AUTO, the format is inferred automatically.	`AUTO`
`valid_graph`	`Optional[Graph]`	Optional. The graph to use for validation. If not provided, the validation set is split from the training graph.	`None`
`num_sampling_hops`	`int`	The number of hops to sample around the seed nodes.	`2`
`sampling_width`	`int`	The number of nodes to sample at each hop.	`15`
`num_layers`	`int`	The number of message passing layers in the GNN.	`2`
`batch_size`	`int`	The batch size to use for training.	`32`
`node_embedding_dim`	`int`	The dimension of the node embeddings.	`128`
`learning_rate`	`float`	The learning rate to use for training.	`0.001`
`cache_valid_dataset`	`bool`	Whether to cache the validation dataset in memory. This can speed up validation if sample generation and normalization are time-consuming, but it will increase memory usage.	`True`
`time_aware`	`bool`	Enables temporal-aware training. If `False` (default), no temporal masking is applied. If `True`, timestamp features are inferred from the schema (via features marked as creation timestamps).	`False`
`message_pooling`	`str`	The pooling method to use for aggregating messages.	`'sum'`
`experimental_preprocess_core_model_config`	`Optional[Callable[[CoreModelConfig], CoreModelConfig]]`	Advanced option. An optional callable to modify the `CoreModelConfig` before it is used to build the core model.	`None`
`cache_normalized_features`	`bool`	If True, pre-compute the normalized features during the preparation stage instead of computing them on the fly in the generator. This option can speed up data generation/training but increases memory consumption.	`False`
`cache_normalized_features_device`	`Literal['host', 'device']`	Specifies the device ("host" for RAM or "device" for GPU/TPU) to store the cached normalized features. Caching features on the same device used for training reduces host-device communication, potentially speeding up training, but increases memory consumption on the device.	`'device'`
`export_metrics_to_xm`	`bool`	If True, metrics from the training and validation steps will be exported to XManager.	`False`
`architecture`	`Union[Architecture, str]`	The architecture of the GNN model.	`DEFAULT_ARCHITECTURE`
`sampling_plan`	`Optional[SamplingPlan]`	An advanced option to provide a custom plan for the sampler. When you use this option, the sampler ignores standard graph sampling arguments and validation checks e.g., num_sampling_hops, sampling_width.	`None`
`diagnostic_dir`	`Optional[str]`	If provided, creates this directory and export to it artefacts that can be useful to understand and debug the model training.	`None`

Returns:

Type	Description
`NodePredictionModel`	A trained `NodePredictionModel` instance.

dgf.learning

dgf.learning.LinkPredictionModel

dgf.learning.Model

__init__

data abstractmethod

describe abstractmethod

name abstractmethod classmethod

save

dgf.learning.NodePredictionModel

evaluate

label_classes

num_label_classes

predict

predict_batch

predict_on_graph_sample_batch

to_tensorflow_function

dgf.learning.load_model

dgf.learning.train_link_model

dgf.learning.train_node_model

`init`

`data` `abstractmethod`

`describe` `abstractmethod`

`name` `abstractmethod` `classmethod`

`save`

`evaluate`

`label_classes`

`num_label_classes`

`predict`

`predict_batch`

`predict_on_graph_sample_batch`

`to_tensorflow_function`