Out of bounds error in TopKPooling with batched graph

[minsikseo-cdl]

Hi, everyone!
I think I found a bug using multiple top k pooling layers with a batched graph. If it's not a bug, please let me know what the problem is.
Following is my modules (abstracted):

class PoolModule(nn.Module):
    def __init__(self, in_dim, out_dim, num_layers, ratio=0.5):
        super().__init__()

        # define GNN before pooling
        self.gnn = gnn.GCN(
            in_dim, out_dim, num_layers,
            act=nn.LeakyReLU(inplace=True), norm=nn.BatchNorm1d(out_dim))

        # Top K Pooling (Gao & Ji, 2019)
        self.pool = gnn.TopKPooling(out_dim, ratio=ratio)

    def forward(self, x, u, edge_index, batch_index):
        out = self.gnn(torch.cat([x, u[batch_index]], dim=-1), edge_index)
        out = self.pool(out, edge_index, batch=batch_index)

        return out[0], out[1], out[3] # x, edge_index, batch_index


class MyModule(nn.Module):
    def __init__(self, args):
        super().__init__()

        # GNN models
        self.homo = HomoModule(**homo_params)
        self.pools = nn.ModuleList()
        for pool_params in pool_param_list:
            self.pools.append(PoolModule(**pool_params))
        self.pred = PredModule(**pred_params)

    def forward(self, in_batch):
        # Homogenization phase
        out_batch = self.homo(in_batch)
        x = out_batch.x
        u = out_batch.u
        edge_index = out_batch.edge_index
        batch_index = out_batch.batch

        # Pooling phase
        for pool in self.pools:
            x, edge_index, batch_index = pool(
                x, u, edge_index, batch_index)
        x = gnn.global_mean_pool(x, batch_index)
        # Prediction phase
        return self.pred(x)

The input batch is a heterogeneous graph and it is homogenized by the first module, namely HomoModule.
During the forward pass in MyModule, a IndexError arises sometimes.

index 517726 is out of bounds for dimension 0 with size 498489
  File "/workspace/topk_pool.py", line 156, in forward
    x = x[perm] * score[perm].view(-1, 1)

(For debugging, I copied the original topk_pool.py in torch_geometric.nn into my workspace.)
Here, shapes of x and score are (498489, 28) and (498489,), respectively.

Basically, is it even possible for perm to exceed the size of x or score?

In HomoModule, to_homogeneous is applied on a heterogeneous graph, such that

out_batch = in_batch.to_homogeneous(
    node_attrs=['x', 'batch'],
    add_node_type=False,
    add_edge_type=False)

Perhaps the problem is that the batch indices in batch are not in order?

If this might be the problem, how can I reorder nodes in the graph?
Or, what is the right way to transfer the batch of a heterogeneous graph to its corresponding homogeneous graph?

Thanks :)

[rusty1s]

Oh, that is interesting. batch needs indeed be in order for TopKPooling to work correctly. However, you are right that to_homogeneous() does not guarantee that for already batched HeteroData objects. I think you have two options:

1. Call to_homogeneous() on each graph individually and use DataLoader to batch them afterwards, e.g.:

class MyTransform():
    def __call__(self, data):
        data = data.to_homogeneuous(node_attrs=['x'], add_node_type=False, add_edge_type=False)
        return data


dataset = MyHeteroDataset(..., transform=MyTransform())
loader = DataLoader(dataset, batch_size=32)

2. Re-order your data object as part of your model:

perm = data.batch.argsort()

data.batch = data.batch[perm]
data.x = data.x[perm]
data.edge_index = perm[data.edge_index]

[minsikseo-cdl]

Thanks @rusty1s !
Reordering the nodes is simpler than I thought! I think I can try the second option and the first method-ish approach. In my HomoModule, there are some MetaLayers to make dimensions of the heterogeneous node and edge attributes same. So I've tried the first-ish method as:

1. Separate HeteroDataBatch to HeteroData list.
2. MetaLayer operation on each HeteroData and apply to_homogeneous().
3. Then batch the homogenized Data and proceed

However, I faced another problem related to the devices. To separate the HeteroDataBatch I tried:

in_graph_list = in_batch.to_data_list()

It gives,

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

.to_data_list() works for a DataBatch but not for HeteroDataBatch. I found the remedy of this but I'm not sure it's right... The problem was that tensors in ._slice_dict and ._inc_dict are not on gpus even though the .cuda() was applied to the data. So I added this before .to_data_list().

device = in_batch.u.device
if device == torch.device('cuda:0'):
    for key, data in in_batch._slice_dict.items():
        if isinstance(data, torch.Tensor):
            in_batch._slice_dict[key] = data.to(device)
        elif isinstance(data, dict):
            for key_, data_ in data.items():
                in_batch._slice_dict[key][key_] = data_.to(device)

    for key, data in in_batch._inc_dict.items():
        if isinstance(data, torch.Tensor):
            in_batch._inc_dict[key] = data.to(device)
        elif isinstance(data, dict):
            for key_, data_ in data.items():
                in_batch._inc_dict[key][key_] = data_.to(device)

But it's kind of ugly to me and the linter also said that it is recommended not to access to the protected members. The code is working so far, though.

In summary, there are two remaining issues:

1. Will the two options below gives identical (or similar) results?
   (1) .to_data_list() -> MetaLayer on each -> .to_homogeneous() -> .from_data_list()
   (2) MetaLayer on batch -> .to_homogeneous() -> reordering
2. Is sending the tensors in ._slice_dict and _inc_dict to a desired device in a manual way okay?

What do you think?

[rusty1s]

Yes, both should give equivalent results, but (2) is probably way faster. Let me also check the root of the issue with to_data_list() and HeteroData. I will keep you posted :)

[rusty1s]

Just fixed the device issues in master, see e6aa4c9.

[minsikseo-cdl]

omg thank you very much for your quick action! :)