Issues when transforming TUDataset

[tindiz]

Hello!

I have an issue and I haven't found a similar one already, so I am opening a new discussion.
As a short intro, I am playing around with a specific class of architectures where I need to transform graphs using EGO policies. However, I still need to keep some information from the original grpah after I do the transformation.

Specifically, I am trying to do something quite similar to what was done in ESAN (https://github.dev/beabevi/ESAN).
Hence, I have a class like this:

class GraphToSubgraph:
    def __init__(self, progress_bar):
        self.progress_bar = progress_bar

    def __call__(self, data):
        if not data.is_undirected():
            raise ValueError("Only undirected graphs are supported!")

        # Create subgraphs and put them in a batch:
        subgraphs = self.to_subgraphs(data)
        batch = Batch.from_data_list(subgraphs)

        # Show progress bar:
        if self.progress_bar is not None:
            next(self.progress_bar)

        # Return data object:
        return SubgraphData(
            x=batch.x,
            edge_index=batch.edge_index,
            edge_attr=batch.edge_attr,
            node2subgraph=batch.batch,
            y=data.y,
            subgraph_idx=batch.subgraph_idx,
            num_subgraphs=len(subgraphs),
            subgraph_node_idx=batch.subgraph_node_idx,
            num_nodes_per_subgraph=data.num_nodes,
            original_edge_index=data.edge_index
        )

    def to_subgraphs(self, data):
        subgraphs = []
        for i in range(data.num_nodes):
            _, _, _, edge_mask = k_hop_subgraph(
                node_idx=i,
                num_hops=3,
                edge_index=data.edge_index,
                relabel_nodes=False,
                num_nodes=data.num_nodes
            )

            subgraph_edge_index = data.edge_index[:, edge_mask]
            subgraph_edge_attr = (
                data.edge_attr[edge_mask]
                if data.edge_attr is not None
                else data.edge_attr
            )

            # Nodes not in the EGO network are isolated nodes.
            subgraphs.append(
                Data(
                    x=data.x,
                    edge_index=subgraph_edge_index,
                    edge_attr=subgraph_edge_attr,
                    subgraph_idx=torch.tensor(i),
                    num_nodes=data.num_nodes,
                    subgraph_node_idx=torch.arange(data.num_nodes),
                )
            )
        return subgraphs

As you mighe notice, there has to be a definition for SubgraphData:

class SubgraphData(Data):
    def __inc__(self, key, value, *args, **kwargs):
        if key == 'original_edge_index':
            return self.num_nodes_per_subgraph
        else:
            return super().__inc__(key, value, *args, **kwargs)

The idea is to keep the original edge index when batching because I will need it later for the model.
When working with the CSL dataset, this works like a charm, I can load the dataset

data_root = '...'
dataset = GNNBenchmarkDataset(root=data_root, name='CSL', transform=None, pre_transform=EgoTransformer(num_hops=4))
train_loader = DataLoader(dataset, batch_size=32, shuffle=True, follow_batch=['subgraph_idx'])

Later when training, we enter into __inc__ and original_edge_index is collated properly:

...
for data in train_loader:  # When stepping in here, __inc__ is called
    # training loop
...

However, this isn't the case when I use the TUDataset. In this case, we have:

data_root = '...'
dataset = TUDataset(root=data_root, name="PROTEINS", transform=None, pre_transform=EgoTransformer(num_hops=4))
train_loader = DataLoader(dataset, batch_size=32, shuffle=True, follow_batch=['subgraph_idx'])

The for loop is the same, but we don't step into __inc__.

I am not sure what I am doing wrong, I would really appreciate some help with this.
Finally, I created a gist so you can try and run the code.

Package version:
pyg==2.4.0

[rusty1s]

Thanks for reporting. You are right that the new data class was not respected when loading the data from disk. Fix is here: #8692

[tindiz]

Thank you, this solves the problem.