如何将3维嵌入张量降维至2维以适配Scikit-learn的K-Means聚类算法？

这个问题我之前也碰到过！核心原因是Scikit-learn的K-Means只接受2维数组输入（形状为(样本数, 特征数)），但你的embeddings_是3维结构——从你给出的数据格式来看，应该是每个样本包含了多个子特征向量（比如文本任务里每个token的embedding），所以得先把每个样本的特征转换成单一向量，才能正常运行K-Means。

下面给你两种最常用的解决方法，结合代码示例说明：

方法一：对每个样本的子向量做池化（推荐） #

这是处理这类多向量样本最常用的方式，通过聚合（均值、最大值、最小值等）把每个样本的多个子向量压缩成一个代表该样本的特征向量。其中均值池化是最通用的选择：

import numpy as np
import sklearn.cluster as cluster

# 先把embeddings转换成numpy数组，确认形状
embeddings = np.array(updating_mapper.embeddings_)
print(f"原始3维数据形状: {embeddings.shape}")  # 输出类似 (N, M, 2)，N是样本数，M是每个样本的子向量数量，2是特征维度

# 对每个样本的所有子向量取均值，得到2维数组 (N, 2)
pooled_embeddings = np.mean(embeddings, axis=1)

# 现在可以正常运行K-Means了
kmeans_labels = cluster.KMeans(n_clusters=10).fit_predict(pooled_embeddings)

如果你想尝试其他聚合方式，比如最大池化，只需要把np.mean换成np.max即可：

pooled_embeddings = np.max(embeddings, axis=1)

方法二：将3维数据展平为高维2维数组 #

如果每个样本的子向量数量不多，你也可以直接把每个样本的所有子向量展平成一个更长的特征向量，这样不会丢失原始信息，但会增加特征维度：

import numpy as np
import sklearn.cluster as cluster

embeddings = np.array(updating_mapper.embeddings_)
# 展平成 (N, M*2) 的2维数组，其中M是每个样本的子向量数量
flattened_embeddings = embeddings.reshape(embeddings.shape[0], -1)

# 运行K-Means
kmeans_labels = cluster.KMeans(n_clusters=10).fit_predict(flattened_embeddings)

额外说明 #

如果你的3维数据结构不是(样本数, 子向量数, 特征数)，而是其他形式（比如(分组数, 样本数, 特征数)），那只需要调整展平/聚合的轴即可。但从你给出的数据格式来看，第一种方法应该完全适用。

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