원문 캡처 · github.com
원문 캡처 · github.com

minibatch-kmeans는 PyTorch 텐서로 미니배치 k-means 군집화를 수행하는 패키지다. CPU와 GPU에서 실행할 수 있으며, 데이터를 배치 단위로 추가 학습해 전체 데이터가 메모리에 들어오지 않아도 된다.

패키지는 Sculley의 Web-Scale K-Means Clustering 알고리즘을 따른다. 각 학습 단계에서 작은 무작위 배치를 뽑아 각 점을 가장 가까운 중심에 할당하고, 중심을 할당된 점 쪽으로 이동시킨다. 중심별로 지금까지 본 점의 수를 세며, 학습률은 그 수의 역수로 설정된다. 이에 따라 각 중심은 할당된 점들의 평균을 유지하고 학습이 진행될수록 이동 폭이 줄어든다.

Python 3.10 이상이 필요하며, 먼저 사용하는 기기에 맞는 PyTorch 2.x를 PyTorch 설치 안내에 따라 설치해야 한다. 이후 패키지를 설치한다.

bash
pip install minibatch-kmeans

fit은 군집 중심을 반환하고 predict는 입력 행마다 하나의 라벨을 반환한다. fit_predict는 두 단계를 한 번에 수행한다. 기본 자료형은 torch.float32이며 입력 텐서의 자료형은 모델의 dtype과 일치해야 한다. 실행마다 같은 결과가 필요하면 시드를 지정한 CPU torch.Generator를 random_state에 전달한다.

python
import torch
from minibatch_kmeans import MiniBatchKMeans

X = torch.randn(5000, 4, dtype=torch.float32)

km = MiniBatchKMeans(n_clusters=8, random_state=torch.Generator().manual_seed(42))
km.fit(X, batch_size=256, max_iter=20)
labels = km.predict(X)

print(labels.shape) # torch.Size([5000])
print(km.cluster_centers.shape) # torch.Size([8, 4])

GPU를 사용하려면 모델과 데이터를 같은 기기로 옮긴다.

python
km = MiniBatchKMeans(n_clusters=8, device=torch.device("cuda"))
labels = km.fit_predict(X.cuda(), batch_size=256)

partial_fit은 전달한 배치의 모든 행으로 중심을 갱신하므로, 데이터가 스트림으로 들어오거나 메모리에 한 번에 담기 어려울 때 사용할 수 있다. 모델은 저장했다가 다시 불러와 학습을 이어갈 수 있다.

python
km = MiniBatchKMeans(n_clusters=8)
for batch in X.split(256):
 km.partial_fit(batch)
labels = km.predict(X)
python
km.save("kmeans.pt")

km2 = MiniBatchKMeans.load("kmeans.pt") # loads on CPU
km_gpu = MiniBatchKMeans.load("kmeans.pt", device=torch.device("cuda")) # loads on a GPU

load는 weights_only 없이 torch.load를 사용하므로 파일에 저장된 코드가 실행될 수 있다. 신뢰하는 파일만 불러와야 한다. 전체 API는 문서에서 확인할 수 있으며, 벤치마크 정보는 평가 안내에 있다. 프로젝트는 MIT 라이선스를 사용한다.