sample_adj for heterograph

[huangtinglin]

SparseTensor supports some very useful graph sampling methods, like graph.sample_adj() and graph.saintgraph().

I wonder if there is a counterpart for heterograph. For example,

from torch_geometric.data import HeteroData

data = HeteroData()

data['paper'].x = ... # [num_papers, num_features_paper]
data['author'].x = ... # [num_authors, num_features_author]
data['institution'].x = ... # [num_institutions, num_features_institution]
data['field_of_study'].x = ... # [num_field, num_features_field]

data['paper', 'cites', 'paper'].edge_index = ... # [2, num_edges_cites]
data['author', 'writes', 'paper'].edge_index = ... # [2, num_edges_writes]
data['author', 'affiliated_with', 'institution'].edge_index = ... # [2, num_edges_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_index = ... # [2, num_edges_topic]

data['paper', 'cites', 'paper'].edge_attr = ... # [num_edges_cites, num_features_cites]
data['author', 'writes', 'paper'].edge_attr = ... # [num_edges_writes, num_features_writes]
data['author', 'affiliated_with', 'institution'].edge_attr = ... # [num_edges_affiliated, num_features_affiliated]
data['paper', 'has_topic', 'field_of_study'].edge_attr = ... # [num_edges_topic, num_features_topic]

# sample an adj starting from the given nodes
src_nodes.= {"papers": [0, 1, 2], "author": [0, 5]}
sampled_data, src_nodes = data.sample_adj(src_nodes, num_neighbors=[5, 5])

I really enjoy using PyG! Looking forward to the reply!

[wsad1]

Will NeighborLoader fit your requirement, you can specify which nodes to start sampling from using input_nodes and it supports HeteroData too.