使用scikit-learn进行k-medoids聚类时小数据集出现内存错误

问题原因 #

你遇到的内存错误是因为KMedoids默认会计算所有样本间的全距离矩阵，当样本量达到426740时，这个矩阵的形状是(426740,426740)，每个元素是8字节的float64，总内存需求约1.33TiB，远远超出常规机器的内存容量。

解决方案 #

1. 使用高效的K-Medoids实现算法 #

sklearn_extra.cluster.KMedoids提供了fastpam1算法，它不需要预计算并存储全距离矩阵，而是在迭代过程中按需计算距离，内存复杂度从O(n²)降到O(n)，适合大样本量场景。修改代码中的KMedoids初始化部分：

kmedoids = KMedoids(
    n_clusters=n_clusters,
    random_state=1,
    algorithm='fastpam1',  # 改用高效算法
    metric='euclidean'     # 可根据需求选择距离度量，比如manhattan
).fit(X)

2. 可选：特征降维（如果特征维度较高） #

如果你的特征维度较多，先通过PCA等方法降维，既能减少计算量，也能降低距离计算的内存开销：

from sklearn.decomposition import PCA

# 先降维到合适的维度，比如保留95%的方差
pca = PCA(n_components=0.95)
X_reduced = pca.fit_transform(X)

# 再用降维后的数据跑KMedoids
kmedoids = KMedoids(n_clusters=25, random_state=1, algorithm='fastpam1').fit(X_reduced)

3. 可选：样本抽样（如果允许近似聚类） #

如果不需要对全样本进行精确聚类，可以先抽取一部分样本训练KMedoids模型，再用模型对剩余样本进行预测：

from sklearn.model_selection import train_test_split

# 随机抽取10%的样本作为训练集
X_train, X_rest = train_test_split(X, test_size=0.9, random_state=1)

# 用训练集训练模型
kmedoids = KMedoids(n_clusters=25, random_state=1, algorithm='fastpam1').fit(X_train)

# 对所有样本预测聚类标签
cluster_labels = kmedoids.predict(X)

修改后的完整代码示例 #

import numpy as np
import pandas as pd

# 导入数据集
dataset = pd.read_csv('dish.csv')
X = dataset.iloc[:, 5:].values

# 处理缺失值
from sklearn.impute import SimpleImputer
imputer = SimpleImputer(missing_values=np.nan, strategy='mean')
X[:, 2:] = imputer.fit_transform(X[:, 2:])

# 特征缩放
from sklearn.preprocessing import StandardScaler
sc = StandardScaler()
X[:, :2] = sc.fit_transform(X[:, :2])

# K-Medoids聚类（使用fastpam1算法）
from sklearn_extra.cluster import KMedoids

n_clusters = 25
kmedoids = KMedoids(
    n_clusters=n_clusters,
    random_state=1,
    algorithm='fastpam1'
).fit(X)

cluster_labels = kmedoids.labels_
dataset['Cluster'] = cluster_labels
dataset.to_csv("clustered_dish.csv", index=False)

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