Full-batch 1-hop subgraph using NeighborLoader

[emduc]

Hi!

I am currently trying to create one full batch that contains all the 1-hop neighbors of my input nodes in the graph. My understanding was that if I only ask NeighborLoader for 1 batch, it should contains only the 1-hop neighbors.

Here is my call:

   batch = list(NeighborLoader(
            graph,
            num_neighbors=[-1], # one layer batch with all neighbors
            time_attr="time",
            input_nodes=table_input.nodes,
            input_time=table_input.time,
            transform=table_input.transform,
            batch_size=graph.num_nodes, # whole graph
            temporal_strategy=cfg.temporal_strategy,
            shuffle=split == "train",
            persistent_workers=False,
        ))

My graph is an HeteroData object. input_nodes of the form ['node_type', Tensor[node IDs]], input_time is a tensor of the same size as input_nodes.

Is there something that I misunderstood so far?

Sadly, the batch.time_dict vectors are much longer than the actual nodes from batch.tf_dict.

I have no idea why that could be the case and any suggestions would be welcome. To tell you a bit more about what I am doing, the graph here is created from a partition of a bigger dataset, the goal is to do distributed training later on, but with a manual partitioning for the purpose of the experiment.

I did consider using k_hop_subgraph, but I am failing to see how that would work with an HeteroData object to get back to a usable HeteroData using the returned values.

Small edit: I do indeed seem to get the 1-hop graph. The problem is with the time vectors

[rusty1s]

What is the error you are seeing? Your understanding is fully correct. I am not sure though what you mean by

Sadly, the batch.time_dict vectors are much longer than the actual nodes from batch.tf_dict.

Both time and tf need to have the shape of the number of nodes of the certain node type in the first dimension.

[emduc]

Hi,

Thanks for your confirmation. By trying out another dataset I discovered that there was a little inconsistency in the timestamps of my partitioning. The behaviour seems indeed correct now. Thanks for your help