KMeans Clustering on GPU with Torch

1

I have a tensor x of shape [32, 10, 128], where:

  • 32 is the batch size,

  • 10 represents nodes,

  • 128 denotes features per
    node.

My objective is to compute node similarities based on their features across each batch input independently. Specifically, I aim to cluster nodes within each input using KMeans clustering on their feature similarities. The end goal is to aggregate nodes that are similar in feature space into clusters and then compute a representative average feature vector for each cluster.

Key requirements:

  • Perform KMeans clustering independently for each input.
  • Utilize Torch to leverage GPU acceleration effectively to execute all the batch inputs parallelly.

Could someone guide me on how to implement this efficiently in Torch?

2

Have you checked a few open-source kmeans implementations, e.g. this one? If so, did it not work for you?

3

I checked your mentioned one, but I think it would not work in my case. Thank you