GPU_KMeans_聚类
更新时间：2026-07-29
简介
cuML GPU K-Means 优先、faiss/sklearn CPU 兜底的向量聚类算子。对一整列 embedding 一次性 fit+predict，输出等长的 cluster_id 列；常与「语义去重（余弦相似度）」串联，实现「先聚 K 簇、再簇内近重复剔除」。
功能描述
• 后端自适应：backend=auto 时按可用性选择 cuml > faiss > sklearn，无 GPU 自动回退 CPU
• 全列聚类：对整列 embedding 一次性 fit+predict，需配合 concurrency=1、batch_size=None 保证同批处理
• 小数据保护：当 n_samples < min_samples_ratio × n_clusters 时自动下调簇数，避免小数据集报错
• 空值稳健：embedding 为 None 的行分配 cluster_id = -1，其余正常聚类
算子参数
输入
| 输入 | 含义 |
|---|---|
| embeddings | embedding 向量列（1 列），元素类型为 list |
输出
| 输出 | 含义 |
|---|---|
| result | int32 cluster_id；embedding 为空的行返回 -1。 |
参数
| 参数名称 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| n_clusters | int | 20000 | 目标簇数 默认值：20000 |
| backend | str | 'auto' | 后端选择 可选值：["auto", "cuml", "faiss", "sklearn"]，auto 时按 cuml>faiss>sklearn 默认值："auto" |
| max_iter | int | 20 | 最大迭代次数 默认值：20 |
| init | str | 'scalable-k-means++' | 初始化策略 可选值：["scalable-k-means++", "random"] 默认值："scalable-k-means++" |
| seed | int | 42 | 随机种子 默认值：42 |
| min_samples_ratio | float | 2.0 | 样本数/簇数下限，不足时自动下调簇数 默认值：2.0 |
| fail_on_error | bool | False | 出错是否抛异常 默认值：False |
调用示例
Python
1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.cluster.gpu_kmeans_cluster import GpuKMeansCluster
10
11if __name__ == "__main__":
12 if os.getenv("DAFT_RUNNER", "native") == "ray":
13 import ray
14 ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15 daft.set_runner_ray()
16 daft.set_execution_config(min_cpu_per_task=0)
17
18 # embeddings：一整列 L2 归一化向量（示例造 3 组明显分离的簇）
19 samples = {
20 "emb": [
21 [0.99, 0.10], [0.98, 0.12], [0.97, 0.14],
22 [0.10, 0.99], [0.12, 0.98], [0.14, 0.97],
23 [-0.70, -0.71], [-0.72, -0.69], [-0.68, -0.73],
24 ]
25 }
26 ds = daft.from_pydict(samples)
27 # 聚类为全列一次性 fit+predict：concurrency=1, batch_size=None
28 ds = ds.with_column(
29 "cluster_id",
30 aihc_udf(
31 GpuKMeansCluster,
32 construct_args={
33 "n_clusters": 3,
34 "backend": "auto",
35 "max_iter": 20,
36 "seed": 42,
37 "min_samples_ratio": 1.0,
38 },
39 num_cpus=1,
40 concurrency=1,
41 batch_size=None,
42 )(col("emb")),
43 )
44 ds.show()评价此篇文章