GSgnnEdgeDataLoaderBase

class graphstorm.dataloading.GSgnnEdgeDataLoaderBase(dataset, target_idx, fanout)

Bases: object

The base dataloader class for edge tasks.

If users want to customize the dataloader for edge prediction tasks they should extend this base class by implementing the special methods __iter__ and __next__.

Parameters

datasetGSgnnEdgeData

The dataset for the edge task.

target_idxdict of Tensors

The target edge IDs.

fanoutlist or dict of lists

The fanout for each GNN layer.

__iter__()

Returns an iterator object

__next__()

Return a mini-batch data for the edge task.

A mini-batch comprises three objects: the input node IDs, the target edges and the subgraph blocks for message passing.

Returns

dict of Tensors : the input node IDs of the mini-batch. DGLGraph : the target edges. list of DGLGraph : the subgraph blocks for message passing.