GSgnnLinkPredictionInferrer

class graphstorm.inference.GSgnnLinkPredictionInferrer(model)

Bases: GSInferrer

Link prediction inferrer.

This is a high-level inferrer wrapper that can be used directly to do link prediction model inference.

Parameters

modelGSgnnNodeModel

The GNN model for node prediction.

property device

The device associated with the inferrer.

property evaluator

Get the evaluator associated with the inferrer.

infer(data, loader, save_embed_path, edge_mask_for_gnn_embeddings='train_mask', use_mini_batch_infer=False, node_id_mapping_file=None, save_embed_format='pytorch')

Do inference

The inference can do two things:

1. (Optional) Evaluate the model performance on a test set if given.

2. Generate node embeddings.

Parameters

data: GSgnnData

The GraphStorm dataset

loaderGSgnnLinkPredictionTestDataLoader

The mini-batch sampler for link prediction task.

save_embed_pathstr

The path where the GNN embeddings will be saved.

edge_mask_for_gnn_embeddingsstr

The mask that indicates the edges used for computing GNN embeddings. By default, the dataloader uses the edges in the training graphs to compute GNN embeddings to avoid information leak for link prediction.

use_mini_batch_inferbool

Whether or not to use mini-batch inference when computing node embedings.

node_id_mapping_file: str

Path to the file storing node id mapping generated by the graph partition algorithm.

save_embed_formatstr

Specify the format of saved embeddings.

setup_device(device)

Set up the device for the inferrer.

The CUDA device is set up based on the local rank.

Parameters

device :

The device for inferrer.

setup_evaluator(evaluator)

Set the evaluator