how to get a cluster vector using in max_pool_x?

[MichaleWong]

❓ Questions & Help

I am tring to implement a graph pooling that is similary to GACNet, I used max_pool_x(),but the parameters cluster vector, I cannot get it efficiently?

I use fps() down sampled the point cloud, and radius() to get each fps point edge_index, try to get cluster vector from the edge_index, but that is very slow, I think it's not a correct way to get cluster vector, do you have some suggestions?

def getCluster(fps_idx, row,col,x):
    r=row.tolist()
    c=col.tolist()
    cluster_tensor = torch.zeros(len(x),dtype=torch.int64)
    icluster=0
    for idx in fps_idx:
        r_idx=[i for i,rx in enumerate(r) if rx==idx.item()]
        idx_neighbors=[c[i] for i in r_idx]
        for idx_n in idx_neighbors:
           cluster_tensor[idx_n]= icluster
        icluster+=1
    return cluster_tensor

how to get the cluster vector correctly?

[rusty1s]

Generally speaking, getting cluster is the hard part :)

Note that the cluster vector in max_pool_x only assigns a single cluster to each node (and I am not sure if that's a good choice here, since then your cluster vector depends on the order of nodes).

A faster implementation may look like this:

cluster = torch.full((x.size(0), ), -1, dtype=torch.long)
cluster[fps_idx] = torch.arange(fps_idx.size(0))
cluster.scatter_(0, row, cluster[col])
cluster[fps_idx] = torch.arange(fps_idx.size(0))
max_pool_x(...)

This implementation assigns clusters to neighbors in a greedy fashion.

If you want to allow multiple clusters per node, implementation is also straightforward (utilizing torch-sparse):

row, col = edge_index
adj_t = SparseTensor(row=row, col=col, sparse_sizes=(x.size(0), x.size(0))).t()
adj_t = adj_t[fps_idx]
x = adj_t.matmul(x, reduce='max')