Is there an equivalent implementation for weight clustering in pytorch as we have in tensorflow : Weight clustering Tesnsorflow 2
If there is not then can someone can someone help me confirming what I have done seems the right thing to do:
from sklearn.cluster import KMeans
# from kmeans_pytorch import kmeans, kmeans_predict

'''Going to go through all the layers --> obtain their weights--> use scikit learn to do the clustering--> replace these weight by their centroids'''
def weight_clustering(model):
  model.to('cpu')
  with torch.no_grad():
    for name, params in model.named_parameters():
      param_shape=list(params.size())
      # print("the shape is ",param_shape)
      weights=torch.flatten(params)
      weights=params.reshape(-1,1)
      kmeans = KMeans(n_clusters=3, random_state=0).fit(weights)

      cluster_centers=torch.from_numpy(kmeans.cluster_centers_)
      for i in range(0,len(kmeans.labels_)):
         if kmeans.labels_[i]==0:
          weights[i]=cluster_centers[0]
         elif kmeans.labels_[i]==1:
          weights[i]=cluster_centers[1]
         elif kmeans.labels_[i]==2:
          weights[i]=cluster_centers[2]
         elif kmeans.labels_[i]==3:
          weights[i]=cluster_centers[3]
         else :
          weights[i]=cluster_centers[4]
         

      reshape_size_tuple=tuple(param_shape)
      weights=weights.reshape(reshape_size_tuple)
      if(params.shape==weights.data.shape):
        params.data=weights.data

Currently, I have tried this very crude way of doing it, but I am not sure this is the right way. Can someone confirm if this is the right way or is there a better/correct way to do the task?