高维PyTorch模型参数KMeans聚类报错问题求助

问题原因 #

你遇到的报错是因为直接收集不同形状的模型参数张量，导致最终输入数组形状不均匀。KMeans要求输入必须是二维数组（样本数 × 特征数），但每个模型的不同参数（比如卷积层权重、偏置）形状各异，直接拼接会产生形状不一致的数组，触发ValueError。

解决方案 #

核心思路是将每个模型的所有参数扁平化拼接成统一长度的一维向量，确保每个样本（模型）对应相同维度的特征向量，再输入KMeans。

步骤1：扁平化每个模型的参数 #

遍历所有模型，将每个模型的所有参数张量扁平化后拼接成一个一维数组：

import torch
import numpy as np
from sklearn.cluster import KMeans

# 假设client_models是同结构的PyTorch模型列表
model_param_vectors = []
for model in client_models:
    # 提取所有参数，逐个扁平化后拼接成单向量
    flat_params = []
    for param in model.parameters():
        flat_params.append(param.flatten().cpu().detach())
    # 拼接成一个一维张量，再转numpy数组
    param_vector = torch.cat(flat_params).numpy()
    model_param_vectors.append(param_vector)

# 转换为二维numpy数组，形状为(模型数量, 总参数数)
param_array = np.array(model_param_vectors)

步骤2：运行KMeans聚类 #

现在可以正常调用KMeans：

kmeans = KMeans(n_clusters=2, random_state=0).fit(param_array)
# 获取每个模型的聚类标签
cluster_labels = kmeans.labels_

可选优化：高维参数降维 #

如果参数维度（30万）过高，KMeans运行效率会很低，还可能出现内存不足问题。可以先用PCA降维，再聚类：

from sklearn.decomposition import PCA

# 降维到合适维度（比如1000维，可根据需求调整）
pca = PCA(n_components=1000, random_state=0)
reduced_params = pca.fit_transform(param_array)

# 基于降维后的数据聚类
kmeans = KMeans(n_clusters=2, random_state=0).fit(reduced_params)

关键注意事项 #

• 必须保证所有client_models是完全相同结构的模型，否则总参数数量不一致，无法拼接成等长的向量。
• 如果是按参数名（state_dict()[k]）单独处理，要确保每个参数名对应的张量形状在所有模型中一致，且最终要将所有参数名对应的张量扁平化后拼接，不能只单独收集某一个参数的张量。

内容的提问来源于stack exchange，提问作者Geetanjli sharma