NeighborLoader failed with pyg-lib

[Xurorx]

Hello PyG team,

Thank you for your great contributions to graph learning science! I recently read that "Once pyg-lib is installed, it will get automatically picked up by PyG, e.g., to accelerate neighborhood sampling routines or to accelerate heterogeneous GNN execution". Thus I want to install pyg-lib to accelerate NeighborLoader sampling.

What I aim to achieve is: given a heterogeneous graph, for specific "edge_type"'s edge_label_index, I want to sample negative target nodes("FILE" node) for each corresponding source node("USER" node). Instead of random negative sampling, I want negative target nodes to be the k-hop neighbors of each source node.

def neighbor_negative_sampling(hetero_data, edge_type, num_hops):
   # get unique node ids for supervision edge's source node
    user_idx = torch.unique(hetero_data[edge_type].edge_label_index[0])
    users = hetero_data['USER'].Id[user_idx]

    loader = NeighborLoader(
        data=data,
        num_neighbors={key: [-1]*num_hops for key in hetero_data.edge_types},
        input_nodes=('USER', users),
        batch_size=1
    )


    hetero_data['USER'].neg_target_nodes = {}

    for sampled_data in loader:
        node_id = sampled_data['USER'].Id[0].item()
        hetero_data['USER'].neg_target_nodes[node_id] = sampled_data['FILE'].Id

    return hetero_data  

However, I got below error message when pyg-lib is installed. Before pyg-lib installation, above code piece worked fine.

 - __main__ - Traceback (most recent call last):
 - __main__ -     return function(*args, **kwargs)
 - __main__ -   File "run.py", line 88, in main
 - __main__ -     neighbor_negative_sampling(data,edge_type,5)
 - __main__ -   File "run.py", line 60, in neighbor_negative_sampling
 - __main__ -     for sampled_data in loader:
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch_geometric/loader/base.py", line 36, in __next__
 - __main__ -     return self.transform_fn(next(self.iterator))
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 628, in __next__
 - __main__ -     data = self._next_data()
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 671, in _next_data
 - __main__ -     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 61, in fetch
 - __main__ -     return self.collate_fn(data)
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch_geometric/loader/node_loader.py", line 112, in collate_fn
 - __main__ -     out = self.node_sampler.sample_from_nodes(input_data)
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch_geometric/sampler/neighbor_sampler.py", line 336, in sample_from_nodes
 - __main__ -     return node_sample(index, self._sample, self.input_type, **kwargs)
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch_geometric/sampler/neighbor_sampler.py", line 376, in node_sample
 - __main__ -     seed_time_dict=seed_time_dict)
 - __main__ -   File "/XXXX/lib/python3.7/site-packages/torch_geometric/sampler/neighbor_sampler.py", line 246, in _sample
 -     row, col, node, edge, batch = out + (None, )
 - ValueError: too many values to unpack (expected 5)

Do you know why "unpack" error appears after installing pyg-lib? Thank you!

[rusty1s]

You need to either install pyg-lib==0.1.0 if you are running on PyG 2.2, or upgrade to PyG 2.3

[Xurorx]

Thank you @rusty1s! I do install PyG 2.2 and pyg-lib==0.2.0 right now. Let me upgrade to PyG2.3.

By the way, is there a function can help us get k-hop_neighbors for a given list of nodes respectively? Right now, I have to iterate over NeighborLoader by batch_size = 1.

[rusty1s]

There exists k_hop_subgraph function. Alternatively, I just added a cleaner interface to NeighborLoader in PyG master, where you can basically do

loader = NeighborLoader(data, ...)
batch = loader(torch.tensor([2]) # get subgraph of node 2.