GSgnnEdgeInferData
- class graphstorm.dataloading.GSgnnEdgeInferData(graph_name, part_config, eval_etypes, label_field=None, node_feat_field=None, edge_feat_field=None, decoder_edge_feat=None)
Bases:
GSgnnEdgeDataEdge prediction inference data
GSgnnEdgeInferData prepares the data for edge prediction inference.
Parameters
- graph_namestr
The graph name
- part_configstr
The path of the partition configuration file.
- eval_etypestuple of str or list of tuples
Target edge types for evaluation
- 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.
- decoder_edge_feat: str or dict of list of str
Edge features used by decoder
Examples
from graphstorm.dataloading import GSgnnEdgeInferData from graphstorm.dataloading import GSgnnEdgeDataLoader ep_data = GSgnnEdgeInferData(graph_name='dummy', part_config=part_config, eval_etypes=[('n1', 'e1', 'n2')], label_field='label', node_feat_field='node_feat', edge_feat_field='edge_feat') ep_dataloader = GSgnnEdgeDataLoader(ep_data, target_idx={"e1":[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(eids, device='cpu')
Get the edge labels
Parameters
- eidsTensor or dict of Tensors
The edge IDs
- devicePytorch device
The device where the returned edge labels are stored.
Returns
dict of Tensors : the returned edge labels.