GSgnnNodeInferData

class graphstorm.dataloading.GSgnnNodeInferData(graph_name, part_config, eval_ntypes, label_field=None, node_feat_field=None, edge_feat_field=None, lm_feat_ntypes=None, lm_feat_etypes=None)

Bases: GSgnnNodeData

Inference data for node tasks

GSgnnNodeInferData prepares the data for node prediction inference.

Parameters

graph_namestr

The graph name

part_configstr

The path of the partition configuration file.

eval_ntypesstr or list of str

Target node types

label_fieldstr

The field for storing labels

node_feat_field: str or dict of list of str

Fields to extract node features. It’s a dict if different node types have different feature names.

edge_feat_fieldstr or dict of list of str

The field of the edge features. It’s a dict if different edge types have different feature names.

lm_feat_ntypeslist of str

The node types that contains text features.

lm_feat_etypeslist of tuples

The edge types that contains text features.

Examples

from graphstorm.dataloading import GSgnnNodeInferData
from graphstorm.dataloading import

np_data = GSgnnNodeInferData(graph_name='dummy', part_config=part_config,
                                eval_ntypes=['n1'], label_field='label',
                                node_feat_field='feat')
np_dataloader = GSgnnNodeDataLoader(np_data, target_idx={'n1':[0]},
                                    fanout=[15, 10], batch_size=128)
get_edge_feats(input_edges, edge_feat_field, device='cpu')

Get the edge features

Parameters

input_edgesTensor or dict of Tensors

The input edge IDs

edge_feat_field: str or dict of [str ..]

The edge data fields that stores the edge features to retrieve

devicePytorch device

The device where the returned edge features are stored.

Returns

dict of Tensors : The returned edge features.

get_labels(nids, device='cpu')

Get the node labels

Parameters

nidsTensor or dict of Tensors

The seed nodes

devicePytorch device

The device where the returned node labels are stored.

Returns

dict of Tensors : the returned node labels.

get_node_feat_size()

Get node feat size using the given node_feat_field

All parameters are coming from this class’s own attributes.

Note: If the self._node_feat_field is None, i.e., not given, the function will return a

dictionary containing all node types in the self.g, and the feature sizes are all 0s. If given the node_feat_field, will return dictionary that only contains given node types.

get_node_feats(input_nodes, device='cpu')

Get the node features

Parameters

input_nodesTensor or dict of Tensors

The input node IDs

devicePytorch device

The device where the returned node features are stored.

Returns

dict of Tensors : The returned node features.

prepare_data(g)

Prepare the testing node set if any

It will setup self._test_idxs, the node indices of the local test set. The test_idxs can be empty.

Arguement

g: DistGraph

The distributed graph.