在 gpu 中使用 pytorch 运行 kmeans 文本聚类以创建超过 1000 个聚类

Question

我正在尝试使用 kmeans-pytorch 实现 kmeans 聚类，但是当我尝试创建超过 10 个聚类时出现内存错误

我的数据集有 7000 条文本记录

这是我的代码片段

import torch

from sklearn.feature_extraction.text import TfidfVectorizer

from kmeans_pytorch import kmeans

text_data = #list of 7000 records 

# Preprocess the data

vectorizer =        TfidfVectorizer(stop_words='english')

X = vectorizer.fit_transform(text_data)

# Convert sparse matrix to PyTorch tensor

X = torch.Tensor(X.toarray())

# Move the data to the GPU

X = X.cuda()

# Run k-means clustering on the GPU
k = 1000

cluster_assignments, centroids = kmeans(X, k,device=torch.device('cuda'))

错误：

RuntimeError: [enforce fail at alloc_cpu.cpp:73] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 41359010000 bytes. Error code 12 (Cannot allocate memory)

Answer 1

需要使用批处理；不幸的是，K-means-pytorch 目前不支持批处理。您可以独立创建批次并找到中心，如 original repo 中定义，或合并，如 ray 和 fast_pytorch_kmenas 中定义。第二种方法将比第一种方法更连贯。