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)

Setup the evaluator

If the evaluator has its own task tracker, just setup the evaluator. But if the evaluator has no task tracker, will use this Trainer’s task tracker to setup the evaluator. When there is no self task tracker, will create a new one by using the given evaluator’s evaluation frequency.