GSgnnEdgePredictionInferrer

class graphstorm.inference.GSgnnEdgePredictionInferrer(model)

Bases: GSInferrer

Edge classification/regression inferrer.

This is a high-level inferrer wrapper that can be used directly to do edge classification/regression 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(loader, save_embed_path, save_prediction_path=None, use_mini_batch_infer=False, node_id_mapping_file=None, edge_id_mapping_file=None, return_proba=True, save_embed_format='pytorch')

Do inference

The inference can do three things:

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

2. Generate node embeddings.

3. Comput inference results for edges with target edge type.

Parameters

loaderGSEdgeDataLoader

The mini-batch sampler for edge prediction task.

save_embed_pathstr

The path where the GNN embeddings will be saved.

save_prediction_pathstr

The path where the prediction results will be saved.

use_mini_batch_inferbool

Whether or not to use mini-batch inference.

node_id_mapping_file: str

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

return_proba: bool

Whether to return all the predictions or the maximum prediction.

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