GSgnnNodePredictionInferrer

class graphstorm.inference.GSgnnNodePredictionInferrer(model)

Bases: GSInferrer

Node classification/regression inferrer.

This is a high-level inferrer wrapper that can be used directly to do node 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, return_proba=True, save_embed_format='pytorch')

Do inference

The inference does three things:

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

2. Generate node embeddings.

3. Comput inference results for nodes with target node type.

Parameters

loaderGSNodeDataLoader

The mini-batch sampler for node 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.