How can I control custom attributes of a Graph in the Data Loader?

[LuisaWerner]

Hi,

I added custom attributes to an OGB data set (PygNodePropPredDataset(name=args.dataset)) that I need for my use case. In my case, this custom attribute is a Tensor with constant values and the size corresponding to the number of edges in the graph.

Later on, I create batches of the data set with NeighborLoader , ClusterLoader and GraphSaintRandomWalkSampler. The edge_index etc. are automatically adapted to the size of the current batch but my added attributes stay at the original size.

How can I implement that the size of my attribute is also adapted to the respective batch?

And how in general can I control how the attributes of the graph as Data object are changed by the loaders?

Thanks!

[wsad1]

NieghborLoader has a transform argument. The transform is applied after the minibatch is created. Maybe you could pass a transform like

class CustomAttribute(BaseTransform):
    def __call__(data):
        data.num_edges = # compute number of edges in data
        data.other_custom_tensor = # set based on data.

[LuisaWerner]

Thanks, that answered my question!

[songsong0425]

Hi, sorry for updating the thread, but I have a similar situation with you. (Please see the detailed contents in here)
In summary, I added the new argument edge_class in the custom Dataset which indicates the source of the original subset, and I want to make a LinkNeighborLoader to get the additional argument edge_class in each batch.
Although I saw the common answers to use torch.utils.data.DataLoader with custom collate_fn, but it seems to change a lot of parent libraries.
How did you solve your problem? Can you share it or any opinion?

[rusty1s]

Will answer you later in #8840, sorry for the delay.

[songsong0425]

@rusty1s Oh, I'm sorry for calling you. I'm just looking for a solution on my own. I don't want to impose on you, so please feel free to reply. 😂