Diffpool Tutorial Does Not Run On Colab

[JosephGatto]

🐛 Describe the bug

I am unable to de-bug this error I receive when running the DiffPool tutorial.

RuntimeError: running_mean should contain 150 elements not 64

I note that in order to get pytorch-geometric to install on Colab in the first place, I did make this adjustment to the install statements in the provided colab notebook - forcing a match in pytorch version.

!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-1.11.0+cu113.html
!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-1.11.0+cu113.html
!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git

Any idea on how to get DiffPool to run on Colab? Thanks!

Environment

• PyG version:
• PyTorch version: 11
• OS: Colab
• Python version:
• CUDA/cuDNN version: 11.3
• How you installed PyTorch and PyG (conda, pip, source): pip
• Any other relevant information (e.g., version of torch-scatter):

[Padarn]

Do you possibly have a link to the collab?

[ivanobat]

You need to modify the DiffPool class so it passes the right size of the x matrix as it is reduced by the pooling. See below how the class looks like now:

class DiffPool(torch.nn.Module):
    def __init__(self):
        super(DiffPool, self).__init__()

        num_nodes_batch = max_nodes
        num_nodes = ceil(0.25 * max_nodes) 
        self.gnn1_pool = GNN(dataset.num_features, 64, num_nodes, num_nodes_batch) 
        self.gnn1_embed = GNN(dataset.num_features, 64, 64, num_nodes_batch)

        num_nodes_batch = num_nodes
        num_nodes = ceil(0.25 * num_nodes)

        self.gnn2_pool = GNN(64, 64, num_nodes, num_nodes_batch)
        self.gnn2_embed = GNN(64, 64, 64, num_nodes_batch, lin=False)

        num_nodes_batch = num_nodes
        self.gnn3_embed = GNN(64, 64, 64, num_nodes_batch, lin=False)

        self.lin1 = torch.nn.Linear(64, 64)
        self.lin2 = torch.nn.Linear(64, dataset.num_classes)

    def forward(self, x, adj, mask=None):
        s = self.gnn1_pool(x, adj, mask)
        x = self.gnn1_embed(x, adj, mask)
        x, adj, l1, e1 = dense_diff_pool(x, adj, s, mask)
        #x_1 = s_0.t() @ z_0
        #adj_1 = s_0.t() @ adj_0 @ s_0

        s = self.gnn2_pool(x, adj)
        x = self.gnn2_embed(x, adj)

        x, adj, l2, e2 = dense_diff_pool(x, adj, s)

        x = self.gnn3_embed(x, adj)

        x = x.mean(dim=1)
        x = F.relu(self.lin1(x))
        x = self.lin2(x)
        return F.log_softmax(x, dim=-1), l1 + l2, e1 + e2