Is there any cluster evaluation method implemented in pyTorch?

1k Views Asked by At

I have found k-means implementation in PyTorch using GPUs which is 30 times faster than CPUs. Is there any method (such as Silhouette score, Dunn's index, ...) implemented preferably in PyTorch that uses GPUs?

1

There are 1 best solutions below

2
Max Schelski On

I implemented the silhouette score in PyTorch based on a numpy implementation from Alexandre Abraham (https://gist.github.com/AlexandreAbraham/5544803):

https://github.com/maxschelski/pytorch-cluster-metrics

After installing the package you can calculate the silhouette score as follows:

from torchclustermetrics import silhouette
score = silhouette.score(X, labels)

With X being the multi-dimensional data (NumPy array or PyTorch tensor; first dimension for samples) and labels being a 1D array of labels for each sample.

I tested the code on PyTorch = 1.10.1 (cuda11.3_cudnn8_0).

In my hands it gave an approximately 30 fold speed up on a GPU compared to the scikit-learn implementation.