Unbatch Data batch from NeighborLoader

[gianmarcoaversanoenx]

Hello

from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='.data', name="Cora")
data = dataset.get(0)

# hyperparameters
n_steps = 2
n_sampled_nodes = 10
batch_size = 2

from torch_geometric.data import LightningNodeData, Data, Batch

datamodule = LightningNodeData(
    data,
    batch_size=batch_size,
    num_neighbors=[n_sampled_nodes] * n_steps,
)

dataloader = datamodule.train_dataloader()

for idx, graph in enumerate(dataloader):
    print(f"({idx}) Data: {graph}\n")
    print(graph.x.size(), graph.edge_index.min(), graph.edge_index.max(), '\n')
    out = query(model, graph)
    break

# OUTPUT WILL BE
# (0) Data: Data(x=[24, 1433], edge_index=[2, 39], y=[24], train_mask=[24], val_mask=[24], test_mask=[24], batch_size=2)

In the code above, I am iterating over a NeighborLoader dataloader that has been initialized with batch_size = 2. So each graph object will be an instance of Data with batch_size=2.

How would I be able to get the 2 graphs (Data object) from graph?

[rusty1s]

In the current version of PyG this is not possible since the neighbors across all seed nodes are shared. There is a new disjoint option in PyG, which will create disjoint subgraphs for each seed node in the batch. This basically lets you re-construct the two data objects.

[gianmarcoaversanoenx]

Thanks for replying!

Where is this disjoint option available? Is there a small code snippet?

[rusty1s]

See here. When using it, it will create a batch vector for each node that maps to its graph.