How to sample custom attributes with NeighborLoader

[dahe-cvl]

Hi,

I would like to sample custom attributes added to a data object.

I created a kNN graph based on an image database and added additional attributes. I was able to generate the graph successfully. The graph has features x, labels y, edge_index as well as edge_attr. Furthermore, I added a kwargs such as "image_path".

During the sampling process with the NeighborLoader class, I am able to receive a batch of nodes with the configured neighborhood. However, I am not getting the "image_paths" from the sampled nodes included in the batch. "image_path" includes all paths from the entire graph.

How can I get only the "image_paths" from the corresponding nodes included in the yielded batch of nodes?

Here is a code snippet:
`
data = dataset[0]

output data --> entire graph

Data(
x=[16767, 512],
edge_index=[2, 502869],
edge_attr=[502869, 1],
y=[16767],
train_mask=[16767],
test_mask=[16767],
kwargs={
img_path=[16767]
}
)

batch_size = 64
sampler_sizes_l = [5, 3, 2]
num_workers = 4

train_loader = NeighborLoader(data,
num_neighbors=sampler_sizes_l,
batch_size=batch_size,
input_nodes=train_idx,
shuffle=True,
num_workers=num_workers,
replace=True,
directed=True,
persistent_workers=False
)

for i, data in enumerate(train_loader):
print(i)
print(data)

---

output: data
Data(
x=[1306, 512],
edge_index=[2, 1944],
edge_attr=[1944, 1],
y=[1306],
train_mask=[1306],
test_mask=[1306],
kwargs={
img_path=[16767], <<--- At this point I should only receive 1306 image_path entries corresponding the the nodes.
},
batch_size=64
)

`

Thanks in advance!
Best regards
Daniel

[dahe-cvl]

Hi, I found a solution.

I add node indices to the entire graph.
data.n_id = torch.arange(data.num_nodes)

With this information, I am able to filter only the image_paths corresponding to the nodes in the batch.

thanks!

[rusty1s]

Amazing. Thank you!

[dahe-cvl]

but give me some time. I think until next week I can share with you a xxx.md file? ok?

[rusty1s]

I don‘t think there is any need to do anything complex. We would just need to add a few lines describing this here: https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/loader/neighbor_loader.py#L169

[dahe-cvl]

ok great, I understand.

[rusty1s]

Added this information to the doc, see here :)