How to use batch_size and num_samples for HGTLoader ?

[LaurentBerder]

Hi,

I'm working with a HeteroData that modelizes bookings, looking roughly like the following for each booking:

Shippers, Consignees and Countries may be connected to different bookings, but a container is only linked to a single booking (though as you can see, a booking may have numerous containers).

I'd like to use the HGTLoader class to get subgraphs containing random bookings, but with all their direct connections (all their containers, their shipper, consignee and countries).

But I don't understand how the num_samples and batch_size arguments work in the class (I'm setting different values for the different node types, getting very different results but none looking like what I want), and the documentation doesn't explain it.

I seem to understand that batch_size is roughly the number of nodes (excluding input_nodes) that we end up with. But as far as the num_samples, I have no idea how to implement it.

Here's an example of results with the following arguments:

loader = HGTLoader(
    data=valid_full,
    batch_size=28,
    num_samples={'booking': [1]*2, 'container': [18]*2, 'shipper': [1]*2, 'consignee': [1]*2, 'country': [3]*2},
    input_nodes='booking'
)
batch = next(iter(loader))
print({k: batch[k]['type'].shape[0] for k in node_types})

{'booking': 29, 'container': 28, 'shipper': 2, 'consignee': 2, 'country': 6}

That looks quite promising, except I've never been able to find a setting that would gie me more than a single container per booking!

How should I approach this?

[rusty1s]

batch_size refers to the number of seed nodes you start sampling from. num_samples refers to the amount of nodes to sample for each node type in each iteration.

It's interesting that the amount of shipper, consignee and country is so low for this. Is every booking connected to one of these?

[rusty1s]

So this is for 3-hops, right? I think this is to be expected. Since you will sample new bookings in the third hop, but then never sample their container nodes since the sampling finishes after three hops. If you are running PyG from master, their is a somewhat hacky way to fix this:

from torch_geometric.sampler import NumNeighbors

N = 2
num_neighbors = NumNeighbors({
  ("container", "rev_has", "booking"): [N] * 4,
}, default = [N] * 3)

loader = NeighborLoader(data, num_neighbors)

which basically means: "Sample three hops, but sample four hops for the container->booking edge type.

[LaurentBerder]

So this is for 3-hops, right?

Yes absolutely. I see what you mean, and did expect some bookings to be incomplete for that reason. But not all of them like it's the case here. The sampler never ever seems to pick more than 1 container node per booking node.

With version 2.2.0, I get cannot import name 'NumNeighbors' from 'torch_geometric.sampler'.
What is the NumNeighbors for, and how is it different from what we're already doing with the dictionnary in NeighborLoader? I have tried before to customize it with something along the line of:

loader = NeighborLoader(
    data,
    # Sample various amounts of neighbors for various edge types for 3 hops
    num_neighbors={('booking', 'has', 'container'): [-1]*3,
                   ('shipper', 'sends', 'booking'): [2]*3,
                   ('consignee', 'receives', 'booking'): [2]*3,
                   ('booking', 'from', 'country'): [2]*3,
                   ('booking', 'to', 'country'): [2]*3,
                   ('shipper', 'from', 'country'): [2]*3,
                   ('consignee', 'from', 'country'): [2]*3,
                   ('container', 'rev_has', 'booking'): [-1]*3,
                   ('booking', 'rev_sends', 'shipper'): [2]*3,
                   ('booking', 'rev_receives', 'consignee'): [2]*3,
                   ('country', 'rev_from', 'booking'): [2]*3,
                   ('country', 'rev_to', 'booking'): [2]*3,
                   ('country', 'rev_from', 'shipper'): [2]*3,
                   ('country', 'rev_from', 'consignee'): [2]*3},
    # Use a batch size of 10 for sampling training nodes of type booking
    batch_size=10,
    input_nodes='booking',
    directed=False
)

With the use of -1 for container-related edges (hoping to force all of them), I have, for the first time, gotten a booking with multiple container nodes:

I then did realize that I have somewhat overestimated the amount of booking nodes linked to multiple container nodes in the data sample I provided...

[rusty1s]

how is it different from what we're already doing with the dictionnary in NeighborLoader?

You are right, it maps to the same, just easier to write and understand :)

[LaurentBerder]

Well, I think we can safely say that we solved it.

I'm now able to get complete subgraphs using NeighborLoader, just like I wanted it.

Thanks for your help, sorry it involved such a lot of your time.

[rusty1s]

Glad that everything works as expected now.