How to use DiffPool after sparse layers?

[j-adamczyk]

I want to use DiffPool as a sort of global pooling, before readout, similar to SAGPool "global" variant (from the SAGPool paper). However, I get errors.

My forward function:

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = self.conv_gnn(x, edge_index)  # node embeddings

        if self.use_diff_pool:
            # adj = to_dense_adj(edge_index, batch)
            s = self.pooling_gnn(x, edge_index)  # cluster assignments
            x, adj, link_pred_loss, entropy = dense_diff_pool(x, edge_index, s)
        else:
            x, edge_index, _, batch, _, _ = self.pooling(x, edge_index, None, batch)

        x = self.readout(x, batch)  # graph embeddings
        x = self.mlp(x)  # (not normalized) class scores

        if self.use_diff_pool:
            return x, link_pred_loss, entropy
        else:
            return x

self.conv_gnn is some ConvGNN with many layers, e.g. GCN or GraphSAGE. self.pooling_gnn is the pooling GNN for DiffPool, e.g. GCNConv(in_channels=conv_out_size, out_channels=num_clusters), creating num_clusters cluster assignments.

I get an error:

  File "C:\Users\jakub\Desktop\Enhanced_GNN_graph_classification\gnn\training.py", line 41, in validation_step
    output = self(data)
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\jakub\Desktop\Enhanced_GNN_graph_classification\gnn\training.py", line 34, in forward
    return self.model(data)
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\jakub\Desktop\Enhanced_GNN_graph_classification\gnn\models\GlobalPoolingGNN.py", line 107, in forward
    x, adj, link_pred_loss, entropy = dense_diff_pool(x, edge_index, s)
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch_geometric\nn\dense\diff_pool.py", line 65, in dense_diff_pool
    out_adj = torch.matmul(torch.matmul(s.transpose(1, 2), adj), s)
RuntimeError: batch1 dim 2 must match batch2 dim 1

If I uncomment adj = to_dense_adj(edge_index, batch) and use adj instead of edge_index, I get another error:

  File "C:\Users\jakub\Desktop\Enhanced_GNN_graph_classification\gnn\training.py", line 34, in forward
    return self.model(data)
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\jakub\Desktop\Enhanced_GNN_graph_classification\gnn\models\GlobalPoolingGNN.py", line 106, in forward
    s = self.pooling_gnn(x, adj)  # cluster assignments
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch\nn\modules\module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch_geometric\nn\conv\gcn_conv.py", line 172, in forward
    edge_index, edge_weight = gcn_norm(  # yapf: disable
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch_geometric\nn\conv\gcn_conv.py", line 58, in gcn_norm
    edge_index, tmp_edge_weight = add_remaining_self_loops(
  File "C:\Users\jakub\anaconda3\envs\Enhanced_GNN_graph_classification\lib\site-packages\torch_geometric\utils\loop.py", line 224, in add_remaining_self_loops
    loop_attr[edge_index[0][inv_mask]] = edge_attr[inv_mask]
IndexError: too many indices for tensor of dimension 1

So can DiffPool be used after sparse layers? If so, how?

[rusty1s]

Note that you need to convert x and s to a "dense batch" as well:

s = self.pooling_gnn(x, edge_index)
x, mask = to_dense_batch(x, batch)
s, _ = to_dense_batch(s, batch)
adj = to_dense_adj(edge_index, batch)
x, adj, link_pred_loss, entropy = dense_diff_pool(x, adj, s, mask)

Let me know if this works :)

[j-adamczyk]

After changing the readout to hand-made self.readout = lambda x: torch.mean(x, dim=1) it works. However, it would be nice to make this work with PyG global pooling layers.

[rusty1s]

Indeed. Fixed my code above. Glad that you resolved it on your own :)

[j-adamczyk]

@rusty1s I am now a bit confused about the readout. global_mean_pool works over dim=0, which makes more sense - it calculates mean of every feature. However, I had to do torch.mean(x, dim=1) for pooling to work. It is the same as in PROTEINS DiffPool example: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/proteins_diff_pool.py. Why are we summing over this axis here? This is also the reason why I could not use global_mean_pool here.

[rusty1s]

In dense mode, the shape of x is [batch_size, num_nodes, num_features], so summing over dim=1 is indeed correct (given you a tensor of shape [batch_size, num_features]).

[j-adamczyk]

Oh, I see, thank you!