BENCH Benchmark 32bit implementation

scikit-learn · Jul 17, 2022 · 2c842bd · 2c842bd
1 parent 31b8b28
commit 2c842bd
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 0 deletions.
diff --git a/asv_benchmarks/asv_scalability.sh b/asv_benchmarks/asv_scalability.sh
@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+#
+# Launch a ASV benchmark on groups of CPU
+# for scalability inspection
+#
+# Results are saved in given folders and files.
+#
+for i in 128 64 32 16 8 4 2 1;
+do
+  last_core=$(($i-1))
+    taskset -c 0-$last_core \
+    asv continuous -b PairwiseDistancesArgKmin \
+    -e main feat/pdr-32bit | tee pairwise_distances_argkmin_asv_${i}_cores.txt
+    cp -R results results_${i}_cores
+done
+
+
+
diff --git a/asv_benchmarks/benchmarks/pairwise_argkmin_estimator.py b/asv_benchmarks/benchmarks/pairwise_argkmin_estimator.py
@@ -0,0 +1,47 @@
+import numpy as np
+
+from .common import Benchmark
+
+from sklearn.metrics import (
+    pairwise_distances_argmin,
+    pairwise_distances_argmin_min
+)
+
+from sklearn.cluster import (
+    AffinityPropagation,
+    Birch,
+    MeanShift,
+    SpectralClustering,
+)
+
+from sklearn.neighbors import NearestNeighbors
+
+from sklearn.manifold import (
+    Isomap,
+    LocallyLinearEmbedding,
+    TSNE,
+)
+
+from sklearn.semi_supervised import (
+    LabelPropagation,
+    LabelSpreading,
+)
+
+class PairwiseDistancesArgKminBenchmark(Benchmark):
+
+    param_names = ["n_train", "n_test", "n_features"]
+    params = [
+        [1000, 10_000, int(1e7)],
+        [1000, 10_000, 100_000],
+        [100],
+    ]
+
+    def setup(self, n_train, n_test, n_features):
+        rng = np.random.RandomState(0)
+        self.X_train = rng.rand(n_train, n_features).astype(np.float32)
+        self.X_test = rng.rand(n_test, n_features).astype(np.float32)
+        self.y_train = rng.randint(low=-1, high=1, size=(n_train,))
+
+    def time_nearest_neighbors(self, n_train, n_test, n_features):
+        est = NearestNeighbors(n_neighbors=10).fit(X=self.X_train)
+        est.kneighbors(self.X_test)