AnswerDotAI / AnswerDotAI/fastkmeans

Torch Tensor and Numpy array

Open
#10 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
104
Forks
8
PR merge metrics
No merged PRs in 30d

Description

Hi 👋, I love fastkmeans which I use a lot,

I'm trying to reduce ram usage of Fast-Plaid which rely on the shoulder of FastKmeans and I would love to avoid a copy of input data. My input embeddings are already instance of torch.Tensor and I convert them to numpy and FastKmeans convert them back to torch.Tensor. Willing to extend the input train method to support `np.ndarray | torch.Tensor`.

```python
def train(self, data: np.ndarray):
"""
Trains (fits) the KMeans model on the given data and sets `self.centroids`. Designed to mimic faiss's `train()` method.

Parameters
----------
data : np.ndarray of shape (n_samples, d), float32
"""
torch.manual_seed(self.seed)
torch.cuda.manual_seed_all(self.seed)
np.random.seed(self.seed)

# Move data to PyTorch CPU Tensor
data_torch = torch.from_numpy(data)
data_norms_torch = (data_torch**2).sum(dim=1)

device = _get_device(self.device)
if device == "cuda" and self.pin_gpu_memory:
data_torch = data_torch.pin_memory()
data_norms_torch = data_norms_torch.pin_memory()

centroids, _ = _kmeans_torch_double_chunked(
data_torch,
data_norms_torch,
k=self.k,
max_iters=self.niter,
tol=self.tol,
device=device,
dtype=self.dtype,
chunk_size_data=self.chunk_size_data,
chunk_size_centroids=self.chunk_size_centroids,
max_points_per_centroid=self.max_points_per_centroid,
verbose=self.verbose,
use_triton=self.use_triton,
)
self.centroids = centroids.numpy()
```

Thank you, FastKmeans is not a bottleneck in FastPlaid, just trying to reduce memory usage one step at a time 😃

I'm making a small PR

Contributor guide

No contributing guide indexed for this repository

Research direction

Start at the train method shown in the issue and trace its call to _kmeans_torch_double_chunked, checking how NumPy data is converted and how centroids are returned. Support the two stated input types without the unwanted conversion copy, then verify training still produces the expected centroids and device behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.