[Neighbor Sampler] Can we specify the edges that sampled from the last layer?

[ShengzheXu]

Hi guys, my question is a little bit complex and let me take an example to show this question:

edge_index_1 = a lot of 100 edges 
edge_index_2 = a lot of 10 edges
combined_edge_index = torch.cat((edge_index_1, edge_index_2), dim=1)

subgraph_loader = NeighborSampler(combined_edge_index, sizes=[50, 50, 50])

100 src_1 ---|
 10  src_2---`--- the tgt_node

Let's look at a node at the last layer, the mini-batch node. Suppose, a target-node that have 100 edges/source-nodes from the edge_index_1 and 10 edges/source-nodes from the edge_index_2. Then we sample 50 neighbors from the total 110 edges.

However, what I focus is actually the 10 edges from edge_index_2 (only consider this case at the last sampled layer). Is there anyway that I can ensure those 10 edges from edge_index_2 can be sampled in the 50-neighbors at the last layer from the NeighborSampler?

[rusty1s]

Mh, interesting, we currently do not provide an interface to allow for this. One workaround to ensure that is to sample all neighbors in the first layer (i.e. sizes = [-1, 50, 50]). Does that help in your case?

[ShengzheXu]

Yeah, I believe this is a very special case. Was wondering if we can save some memory :)

I turned to run like [-1, 50, 50]. Highly appreciate your suggestion!