FIX Number of threads in KMeans should not be bigger than number of chunks (scikit-learn#17210)

jeremiedbb · adrinjalali · commit 15716dabf8a1 · 2020-05-19T10:00:50.000+02:00
* num threads not bigger than num chunks

* what's new
diff --git a/doc/whats_new/v0.23.rst b/doc/whats_new/v0.23.rst
@@ -19,6 +19,10 @@ Changelog
   provided by the user were modified in place. :pr:`17204` by
   :user:`Jeremie du Boisberranger <jeremiedbb>`.
 
+- |Efficiency| :class:`cluster.KMeans` cannot spawn idle threads any more for
+  very small datasets. :pr:`17210` by
+  :user:`Jeremie du Boisberranger <jeremiedbb>`.
+
 Miscellaneous
 .............
 
diff --git a/sklearn/cluster/_k_means_elkan.pyx b/sklearn/cluster/_k_means_elkan.pyx
@@ -284,6 +284,9 @@ def _elkan_iter_chunked_dense(
     # count remainder chunk in total number of chunks
     n_chunks += n_samples != n_chunks * n_samples_chunk
 
+    # number of threads should not be bigger than number of chunks
+    n_threads = min(n_threads, n_chunks)
+
     if update_centers:
         memset(&centers_new[0, 0], 0, n_clusters * n_features * sizeof(floating))
         memset(&weight_in_clusters[0], 0, n_clusters * sizeof(floating))
@@ -515,6 +518,9 @@ def _elkan_iter_chunked_sparse(
     # count remainder chunk in total number of chunks
     n_chunks += n_samples != n_chunks * n_samples_chunk
 
+    # number of threads should not be bigger than number of chunks
+    n_threads = min(n_threads, n_chunks)
+
     if update_centers:
         memset(&centers_new[0, 0], 0, n_clusters * n_features * sizeof(floating))
         memset(&weight_in_clusters[0], 0, n_clusters * sizeof(floating))
diff --git a/sklearn/cluster/_k_means_lloyd.pyx b/sklearn/cluster/_k_means_lloyd.pyx
@@ -120,6 +120,9 @@ def _lloyd_iter_chunked_dense(
     # count remainder chunk in total number of chunks
     n_chunks += n_samples != n_chunks * n_samples_chunk
 
+    # number of threads should not be bigger than number of chunks
+    n_threads = min(n_threads, n_chunks)
+
     if update_centers:
         memset(&centers_new[0, 0], 0, n_clusters * n_features * sizeof(floating))
         memset(&weight_in_clusters[0], 0, n_clusters * sizeof(floating))
@@ -314,6 +317,9 @@ def _lloyd_iter_chunked_sparse(
     # count remainder chunk in total number of chunks
     n_chunks += n_samples != n_chunks * n_samples_chunk
 
+    # number of threads should not be bigger than number of chunks
+    n_threads = min(n_threads, n_chunks)
+
     if update_centers:
         memset(&centers_new[0, 0], 0, n_clusters * n_features * sizeof(floating))
         memset(&weight_in_clusters[0], 0, n_clusters * sizeof(floating))
