Question about avg_pool_x lost batch information

[YunYunY]

Hi,

Background: I am implementing spatial average pooling after eq. 5 of paper
From my understanding, avg_pool_x find the average of nodes among same cluster. Is it averating in the same cluster of all batch?
Suppose my input graph has different sizes. Batch_size is 5, each node belong to one of the 4 clusters. The cluster of node is like

tensor([0, 1, 1, 2, 2, 3, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0,
        1, 1, 2, 2, 3, 3, 0, 1, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3], device='cuda:0')

If I do

   def forward(self, data):
        batch = data.batch
        x, edge_index, edge_weight, cluster = data.x, data.edge_index, data.edge_attr, data.partition        
        x = F.relu(self.conv1(x, edge_index, edge_weight))
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index, edge_weight)
        x, batch = avg_pool_x(cluster, x, data.batch)
        print(x.shape)

I got x.shape is torch.Size([4, 87]), 87 is num_classes. I lost the original batchwise information here right? How could I average for each cluster in the same batch? Say the batch_size is 5, I'm looking for a way to get batch_size * n_cluster * num_classes so that I keep the cluster information like T on page 10873 right bottom of the paper.

Thank you very much for your time.

---

I found one related discussion here. Regarding the answer for the issue, shouldn't we use the following method to create clusters for different batch?

  batch = data.batch
  print(batch)
  print(cluster)
  cluster_max = cluster.max().item() + 1
  batch_size = cluster.shape[0]
  new_cluster = batch * cluster_max + cluster
  print(new_cluster)

This will print

batch: tensor([0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2,
        2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4], device='cuda:0')
cluster: tensor([0, 1, 1, 2, 2, 3, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0,
        1, 1, 2, 2, 3, 3, 0, 1, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3], device='cuda:0')
new_cluster: tensor([ 0,  1,  1,  2,  2,  3,  4,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  6,
         7,  7,  7,  7,  8,  8,  9,  9, 10, 10, 11, 11, 12, 13, 13, 14, 15, 16,
        16, 17, 17, 18, 18, 19], device='cuda:0')

If creating the new_cluster with this we will get

new_cluster = torch.arange(batch_size).cuda() * cluster_max + cluster
print(new_cluster)
new_cluster: tensor([  0,   5,   9,  14,  18,  23,  24,  28,  32,  36,  41,  45,  49,  53,
         58,  62,  66,  70,  75,  79,  83,  87,  88,  92,  97, 101, 106, 110,
        115, 119, 120, 125, 129, 134, 139, 140, 144, 149, 153, 158, 162, 167],
       device='cuda:0')

Every node would be assigned to a different cluster.

I applied

cluster = batch * cluster_max + cluster   
x, _ = avg_pool_x(x, cluster, batch=data.batch)

But the error shows:

  result = self.forward(*input, **kwargs)
  File "/model/gcn.py", line 58, in forward
    x, _ = avg_pool_x(x, cluster, batch=data.batch)
  File "/anaconda3/envs/graph/lib/python3.6/site-packages/torch_geometric/nn/pool/avg_pool.py", line 37, in avg_pool_x
    cluster, perm = consecutive_cluster(cluster)
  File "/anaconda3/envs/graph/lib/python3.6/site-packages/torch_geometric/nn/pool/consecutive.py", line 5, in consecutive_cluster
    unique, inv = torch.unique(src, sorted=True, return_inverse=True)
  File "/anaconda3/envs/graph/lib/python3.6/site-packages/torch/_jit_internal.py", line 265, in fn
    return if_true(*args, **kwargs)
  File "/anaconda3/envs/graph/lib/python3.6/site-packages/torch/_jit_internal.py", line 267, in fn
    return if_false(*args, **kwargs)
  File "/anaconda3/envs/graph/lib/python3.6/site-packages/torch/functional.py", line 778, in _return_inverse
    output, inverse_indices, _ = _unique_impl(input, sorted, return_inverse, return_counts, dim)
  File "/anaconda3/envs/graph/lib/python3.6/site-packages/torch/functional.py", line 682, in _unique_impl
    return_counts=return_counts,
RuntimeError: isDifferentiableType(variable.scalar_type()) INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1607370116979/work/torch/csrc/autograd/functions/utils.h":64, please report a bug to PyTorch.

Is this a bug or I misunderstood something?
Thank you for your help.

[rusty1s]

You are right that you need to modify cluster to account for different examples in a batch. For _avg_pool_x, you need to first pass in the cluster, and then the value you want to pool:

import torch
from torch_geometric.nn import avg_pool_x

batch = torch.tensor([
    0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2,
    2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4
])
x = torch.randn(batch.size(0), 128)
cluster = torch.tensor([
    0, 1, 1, 2, 2, 3, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0, 1,
    1, 2, 2, 3, 3, 0, 1, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3
])

cluster_max = cluster.max().item() + 1
new_cluster = batch * cluster_max + cluster

out, _ = avg_pool_x(new_cluster, x, batch)