Pytorch 中的权重聚类

Question

我一直在为我的项目尝试在 PyTorch 中进行权重聚类。 为此，我使用了以下非常基本的代码：

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)
      params.data=weights.data

但是，这将单独保存权重，因此不会减少 model 大小。 我想使用和访问相同的存储（类似于以下链接： PyTorch Views ）。 据我所知，与 TensorFlow 相比，我相信 PyTorch 还没有任何集群机制。

我曾尝试在上述代码中使用view ，但未能成功获得所需的结果。有人可以指导/帮助我吗？