[FEA] Add Batching to KMeans#1886
Merged
rapids-bot[bot] merged 170 commits intorapidsai:release/26.04from Mar 31, 2026
Merged
Conversation
…into minibatch-kmeans
…into minibatch-kmeans
…into minibatch-kmeans
lowener
requested changes
Mar 26, 2026
Contributor
Author
|
Taking into consideration these comments: #1886 (comment) and #1886 (comment) I have removed the host side predict function and put them into #1962 with DO NOT MERGE status to track those changes. |
tarang-jain
commented
Mar 28, 2026
|
|
||
| assert np.allclose( | ||
| inertia_regular, inertia_batched, rtol=1e-3, atol=1e-3 | ||
| ), f"max diff: {np.max(np.abs(inertia_regular - inertia_batched))}" |
Contributor
Author
There was a problem hiding this comment.
Added this inertia check to the python side to evaluate that the inertia calculation is correct in the batched fit function.
tarang-jain
commented
Mar 28, 2026
| raft::host_vector_view<IdxT, IdxT> labels, | ||
| bool normalize_weight, | ||
| raft::host_scalar_view<T> inertia) | ||
| { |
Contributor
Author
There was a problem hiding this comment.
I have still left this function here in the detail namespace because I think there is some utility to it. Its not being used anywhere and tests and public API have been moved to #1962
Contributor
Author
aamijar
approved these changes
Mar 30, 2026
lowener
approved these changes
Mar 31, 2026
Contributor
lowener
left a comment
lowener
left a comment
There was a problem hiding this comment.
LGTM, only a minor remark
jinsolp
approved these changes
Mar 31, 2026
Contributor
jinsolp
left a comment
jinsolp
left a comment
There was a problem hiding this comment.
Thanks for the great work @tarang-jain LGTM!
Contributor
Author
|
/merge |
benfred
approved these changes
Mar 31, 2026
github-project-automation
bot
moved this from In Progress
to Done
in Unstructured Data Processing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Merge after #1880
This PR adds support for streaming out of core (dataset on host) kmeans clustering. The idea is simple:
Batched accumulation of centroid updates: Data is processed in batches and batch-wise means and cluster counts are accumulated until all the batches i.e., the full dataset pass has completed.
This PR just brings a batch-size parameter to load and compute cluster assignments and (weighted) centroid adjustments on batches of the dataset. The final centroid 'updates' i.e. a single kmeans iteration only completes when all these accumulated sums are averaged once the whole dataset pass has completed.