perf(model): Replace Sklearn with optimized PyTorch KMeans in CFA by andersendsa · Pull Request #3760 · open-edge-platform/anomalib

andersendsa · 2026-09-01T20:00:52Z

📝 Description

Addresses TODO in src/anomalib/models/image/cfa/torch_model.py to use a PyTorch KMeans class instead of Sklearn's version which requires moving data to the CPU. - Uses anomalib.models.components.cluster.kmeans.KMeans and optimizes it using PyTorch's vectorized bincount and index_add_ operations instead of python loops with boolean indexing, providing significant performance speedups.

✨ Changes

Select what type of change your PR is:

✅ Checklist

Before you submit your pull request, please make sure you have completed the following steps:

📚 I have made the necessary updates to the documentation (if applicable).
🧪 I have written tests that support my changes and prove that my fix is effective or my feature works (if applicable).
🏷️ My PR title follows conventional commit format.

For more information about code review checklists, see the Code Review Checklist.

Copilot

🟡 Changes recommended

KMeans currently lacks required n_clusters validation and uses CPU-created centroid indices (risking runtime failure / device sync), and touched-file copyright headers need updating to include 2026.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR replaces the CFA model’s use of scikit-learn KMeans (CPU-bound) with anomalib’s PyTorch KMeans implementation and improves KMeans centroid updates using vectorized PyTorch operations to reduce Python-loop overhead and avoid unnecessary device transfers.

Changes:

Swap CFA centroid initialization from sklearn.cluster.KMeans to anomalib.models.components.cluster.kmeans.KMeans.
Optimize KMeans centroid recomputation using torch.bincount + index_add_ instead of per-cluster boolean masking loops.

File summaries

File	Description
src/anomalib/models/image/cfa/torch_model.py	Switch CFA memory-bank clustering from sklearn KMeans (CPU) to anomalib’s PyTorch KMeans.
src/anomalib/models/components/cluster/kmeans.py	Vectorize centroid updates and adjust centroid initialization behavior.

Review details

Files reviewed: 2/2 changed files
Comments generated: 3
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

ashwinvaidya17

Thanks for updating this. I have left a very minor comment. Can you also have a look at the copilot comments as well. They are quite minor as well

ashwinvaidya17 · 2026-09-02T11:49:08Z

Also, can you update the PR title

Copilot

🟡 Changes recommended

The new KMeans implementation currently lacks convergence-based early exit (critical with max_iter=3000) and CFA’s clustered memory bank dtype now differs from the prior explicit float32 behavior.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

src/anomalib/models/components/cluster/kmeans.py:101

This vectorized centroid-update logic is performance-critical and can materially change clustering behavior (especially around empty clusters and convergence). There are unit tests for GaussianMixture that rely on KMeans indirectly, but there are no direct tests for KMeans itself; adding a small deterministic unit test (synthetic blobs, fixed seed, asserts on shapes + stable clustering) would help prevent regressions in this optimized path.

Files reviewed: 2/2 changed files
Comments generated: 2
Review effort level: Lite

Copilot

🟡 Changes recommended

The PyTorch KMeans implementation will fail on accelerator devices due to CPU-created index tensors, and the new/updated clustering behavior lacks direct unit test coverage.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

src/anomalib/models/components/cluster/kmeans.py:85

centroid_indices is created on the default device (CPU). If inputs is on CUDA/XPU, inputs[centroid_indices] will error because the index tensor must be on the same device as the indexed tensor. This becomes likely now that CFA calls KMeans on-device to avoid CPU transfers.

        # Initialize centroids randomly from the data points
        centroid_indices = torch.randint(0, batch_size, (self.n_clusters,))
        self.cluster_centers_ = inputs[centroid_indices].clone()

Files reviewed: 2/2 changed files
Comments generated: 1
Review effort level: Lite

Copilot

🟡 Changes recommended

The new vectorized KMeans update path can produce numerically incorrect centroids under mixed precision (float16 counts) and the CFA migration should keep the prior float32 behavior and add regression coverage.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

src/anomalib/models/image/cfa/torch_model.py:248

cluster_centers are now kept in whatever dtype self.memory_bank has; previously the sklearn path forced float32. Under mixed precision/autocast this can make the memory bank float16, which can reduce numerical stability in compute_distance and change outputs compared to the prior implementation. Consider running KMeans in float32 and storing the resulting centers as float32 (matching the old behavior).

            _, cluster_centers = k_means.fit(self.memory_bank)
            self.memory_bank = cluster_centers.detach().to(device)

Files reviewed: 2/2 changed files
Comments generated: 2
Review effort level: Lite

+            counts = torch.bincount(self.labels_, minlength=self.n_clusters).to(inputs.dtype)
+            new_centers = torch.zeros_like(self.cluster_centers_)
+            new_centers.index_add_(0, self.labels_, inputs)
+
+            valid_mask = counts > 0
+            self.cluster_centers_[valid_mask] = new_centers[valid_mask] / counts[valid_mask].unsqueeze(1)


        if self.gamma_c > 1:
-            # TODO(samet-akcay): Create PyTorch KMeans class.
-            # CVS-122673
            k_means = KMeans(


Added unit tests for KMeans to assert that fit returns expected shapes and dtypes, and that predict is consistent with fitted centers on a synthetic dataset as requested in the PR comments.

andersendsa · 2026-09-09T13:53:28Z

Hi @ashwinvaidya17 I have added unit tests could you pls let me know if this pr needs any changes or if it is good to merge

Copilot

🟡 Changes recommended

The updated PyTorch KMeans has correctness issues for GPU usage and can return labels inconsistent with the final cluster centers unless adjusted.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

src/anomalib/models/components/cluster/kmeans.py:85

centroid_indices is created on the default device (CPU). If inputs is on CUDA (common for CFA/GMM), indexing inputs[centroid_indices] can raise a device-mismatch error. Create the indices on inputs.device to ensure KMeans works on GPU.

        # Initialize centroids randomly from the data points
        centroid_indices = torch.randint(0, batch_size, (self.n_clusters,))
        self.cluster_centers_ = inputs[centroid_indices].clone()

Files reviewed: 3/3 changed files
Comments generated: 2
Review effort level: Lite

ashwinvaidya17 · 2026-09-10T12:45:51Z

@andersendsa I know continuous copilot comments can be annoying but can you check if they make sense and update the PR if necessary?

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

andersendsa · 2026-09-17T06:27:51Z

@andersendsa I know continuous copilot comments can be annoying but can you check if they make sense and update the PR if necessary?

Hi @ashwinvaidya17 sorry for the delay I have implemented the changes required and resolved them pls let me know if this pr needs anymore changes or if it is good to merge

ashwinvaidya17

Thanks for the efforts. Just one minor comment. I think we should add early exit once the labels stop changing. Otherwise we might keep it running till max_iter

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

andersendsa · 2026-09-22T20:22:00Z

Thanks for the efforts. Just one minor comment. I think we should add early exit once the labels stop changing. Otherwise we might keep it running till max_iter

Hi @ashwinvaidya17 I have done the changes pls let me know if pr needs anymore changes or if it is good to merge

ashwinvaidya17 · 2026-09-23T11:30:31Z

@andersendsa can you fix the prek issues? I think we are good to merge after that.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

andersendsa · 2026-09-24T20:37:06Z

@andersendsa can you fix the prek issues? I think we are good to merge after that.

Hi @ashwinvaidya17 I have fixed the prek issues and This pr is good to merge

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Replace Sklearn with optimized PyTorch KMeans in CFA

24108ca

Copilot AI lite review requested due to automatic review settings September 1, 2026 20:00

andersendsa requested review from ashwinvaidya17, rajeshgangireddy and samet-akcay as code owners September 1, 2026 20:00

Copilot started reviewing on behalf of andersendsa September 1, 2026 20:01 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Comment thread src/anomalib/models/components/cluster/kmeans.py

Comment thread src/anomalib/models/components/cluster/kmeans.py Outdated

Comment thread src/anomalib/models/image/cfa/torch_model.py

ashwinvaidya17 requested changes Sep 2, 2026

View reviewed changes

Comment thread src/anomalib/models/image/cfa/torch_model.py Outdated

minor changes

ec641dd

Copilot AI review requested due to automatic review settings September 2, 2026 14:03

Copilot started reviewing on behalf of andersendsa September 2, 2026 14:04 View session

andersendsa changed the title ~~Replace Sklearn with optimized PyTorch KMeans in CFA~~ perf(model) : Replace Sklearn with optimized PyTorch KMeans in CFA Sep 2, 2026

Merge branch 'main' into kmeans-optimization

c5f8ca0

andersendsa requested a review from ashwinvaidya17 September 2, 2026 14:06

Copilot AI reviewed Sep 2, 2026

View reviewed changes

Comment thread src/anomalib/models/components/cluster/kmeans.py Outdated

Comment thread src/anomalib/models/image/cfa/torch_model.py

Copilot AI review requested due to automatic review settings September 2, 2026 14:08

Copilot started reviewing on behalf of andersendsa September 2, 2026 14:09 View session

Copilot AI reviewed Sep 2, 2026

View reviewed changes

Comment thread src/anomalib/models/components/cluster/kmeans.py Outdated

ashwinvaidya17 changed the title ~~perf(model) : Replace Sklearn with optimized PyTorch KMeans in CFA~~ perf(model): Replace Sklearn with optimized PyTorch KMeans in CFA Sep 3, 2026

Merge branch 'main' into kmeans-optimization

3fb6fd0

Copilot AI review requested due to automatic review settings September 8, 2026 15:44

Copilot started reviewing on behalf of andersendsa September 8, 2026 15:44 View session

Copilot AI reviewed Sep 8, 2026

View reviewed changes

andersendsa and others added 2 commits September 9, 2026 13:47

test: add focused tests for KMeans clustering

5b18d12

Added unit tests for KMeans to assert that fit returns expected shapes and dtypes, and that predict is consistent with fitted centers on a synthetic dataset as requested in the PR comments.

adding tests

fa8dc6e

Copilot AI review requested due to automatic review settings September 9, 2026 13:51

Copilot started reviewing on behalf of andersendsa September 9, 2026 13:52 View session

Copilot AI reviewed Sep 9, 2026

View reviewed changes

Comment thread src/anomalib/models/components/cluster/kmeans.py

Comment thread src/anomalib/models/components/cluster/kmeans.py Outdated

implementing suggestions

5e2a690

Copilot AI review requested due to automatic review settings September 17, 2026 06:21

Copilot AI reviewed Sep 17, 2026

Copilot AI review requested due to automatic review settings September 17, 2026 06:27

Copilot AI reviewed Sep 17, 2026

ashwinvaidya17 reviewed Sep 18, 2026

View reviewed changes

adding early exit

e3a8dec

Copilot AI review requested due to automatic review settings September 18, 2026 12:59

andersendsa force-pushed the kmeans-optimization branch from 34058e2 to e3a8dec Compare September 18, 2026 12:59

Copilot AI reviewed Sep 18, 2026

andersendsa requested a review from ashwinvaidya17 September 18, 2026 13:00

format fix

256fa99

Copilot AI review requested due to automatic review settings September 24, 2026 20:34

Copilot AI reviewed Sep 24, 2026

Merge branch 'main' into kmeans-optimization

e2a43a6

Copilot AI review requested due to automatic review settings September 29, 2026 09:05

Copilot AI reviewed Sep 29, 2026

ashwinvaidya17 approved these changes Sep 29, 2026

View reviewed changes

ashwinvaidya17 merged commit 605b531 into open-edge-platform:main Sep 29, 2026
37 checks passed

andersendsa deleted the kmeans-optimization branch September 29, 2026 12:50

Conversation

andersendsa commented Sep 1, 2026

📝 Description

✨ Changes

✅ Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ashwinvaidya17 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ashwinvaidya17 commented Sep 2, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

andersendsa commented Sep 9, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (1)

Uh oh!

Uh oh!

Uh oh!

ashwinvaidya17 commented Sep 10, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

andersendsa commented Sep 17, 2026

Uh oh!

ashwinvaidya17 left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

andersendsa commented Sep 22, 2026

Uh oh!

ashwinvaidya17 commented Sep 23, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

andersendsa commented Sep 24, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone