AnswerDotAI / AnswerDotAI/fastkmeans
Torch Tensor and Numpy array
- 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