dataloader excluding other external features added to datasets.data

[sidhantls]

The dataloader (torch_geometric.data.DataLoader) excludes other variables (alternate features) that are added to dataset.data. For example, dataset.data.ext_features is excluded when generating the batch. Is there any way around this?

I am using a dataset from PG torch_geometric.datasets- WikiCS which contains only 1 graph. I include external features, that are not part of the given features in a given dataset- I add the external features to dataset.data.ext_features, and set up forward pass such that I use this.
however, when the data loader generates the batch, this variable is excluded and only the regular edge_index, stopping_mask. test_mask, train_mask, val_mask, x, y are present. Not sure why the dataloader isn't including ext_features in the batch, since it is a variable in dataset.data. Is there any way around this?

[rusty1s]

Modifying a dataset via dataset.data is not really recommended since this is only meant as an internal storage layer. For this to work, you also need to set the dataset.slices object, and reset caching:

import torch
from torch_geometric.datasets import WikiCS

dataset = WikiCS('/tmp/Wiki')
data = dataset[0]
dataset.data.ext_features = torch.randn(data.num_nodes, 128)
dataset.slices['ext_features'] = torch.tensor([0, data.num_nodes])
dataset._data_list = None  # Reset caching.
print(dataset[0])

A better way to do this is via transforms:

def add_features(data):
    data.ext_features = ... # Add your features

dataset = WikiCS('/tmp/Wiki', transform=add_features)

[sidhantls]

Just fyi, the first code sample shared doesn't seem to work for this case

print(dataset[0]) does not show the ext_feature variable in dataset.data. Although, print(dataset.data) shows that it is registered. After getting the batch from dataloaders, batch.ext_feature isn't present

import torch
from torch_geometric.datasets import WikiCS
from torch_geometric.data import DataLoader

dataset = WikiCS('/tmp/Wiki')
data = dataset[0]
dataset.data.ext_features = torch.randn(data.num_nodes, 128)
dataset.slices['ext_features'] = torch.tensor([0, data.num_nodes])
dataset._data_list = None  # Reset caching.
print(dataset[0])

loader=DataLoader(dataset)

for batch in loader:
    break

# passes
assert hasattr(batch, 'x')   
# fails 
assert hasattr(batch, 'ext_features')

[sidhantls]

Also the second code sample also doesnt seem to work. the add_features function isn't even being called:

def add_features(data):
    data.ext_features = torch.randn(data.num_nodes, 128)
    raise ValueError()
   
# no error raised
dataset = WikiCS(data_dir, transform=add_features)

[rusty1s]

Mh, the first script does not fail for me. For the second script, note that you need to access a data object of your dataset in order for the transform to be called.

def add_features(data):
    data.ext_features = torch.randn(data.num_nodes, 128)
    raise ValueError()
   
# no error raised
dataset = WikiCS(data_dir, transform=add_features)

loader = DataLoader(dataset)

for batch in loader:
    break