143 lines
4.3 KiB
Python
143 lines
4.3 KiB
Python
from collections import defaultdict
|
|
from time import time
|
|
|
|
import numpy as np
|
|
from numpy import random as nr
|
|
|
|
from sklearn.cluster import KMeans, MiniBatchKMeans
|
|
|
|
|
|
def compute_bench(samples_range, features_range):
|
|
|
|
it = 0
|
|
results = defaultdict(lambda: [])
|
|
chunk = 100
|
|
|
|
max_it = len(samples_range) * len(features_range)
|
|
for n_samples in samples_range:
|
|
for n_features in features_range:
|
|
it += 1
|
|
print("==============================")
|
|
print("Iteration %03d of %03d" % (it, max_it))
|
|
print("==============================")
|
|
print()
|
|
data = nr.randint(-50, 51, (n_samples, n_features))
|
|
|
|
print("K-Means")
|
|
tstart = time()
|
|
kmeans = KMeans(init="k-means++", n_clusters=10).fit(data)
|
|
|
|
delta = time() - tstart
|
|
print("Speed: %0.3fs" % delta)
|
|
print("Inertia: %0.5f" % kmeans.inertia_)
|
|
print()
|
|
|
|
results["kmeans_speed"].append(delta)
|
|
results["kmeans_quality"].append(kmeans.inertia_)
|
|
|
|
print("Fast K-Means")
|
|
# let's prepare the data in small chunks
|
|
mbkmeans = MiniBatchKMeans(
|
|
init="k-means++", n_clusters=10, batch_size=chunk
|
|
)
|
|
tstart = time()
|
|
mbkmeans.fit(data)
|
|
delta = time() - tstart
|
|
print("Speed: %0.3fs" % delta)
|
|
print("Inertia: %f" % mbkmeans.inertia_)
|
|
print()
|
|
print()
|
|
|
|
results["MiniBatchKMeans Speed"].append(delta)
|
|
results["MiniBatchKMeans Quality"].append(mbkmeans.inertia_)
|
|
|
|
return results
|
|
|
|
|
|
def compute_bench_2(chunks):
|
|
results = defaultdict(lambda: [])
|
|
n_features = 50000
|
|
means = np.array(
|
|
[
|
|
[1, 1],
|
|
[-1, -1],
|
|
[1, -1],
|
|
[-1, 1],
|
|
[0.5, 0.5],
|
|
[0.75, -0.5],
|
|
[-1, 0.75],
|
|
[1, 0],
|
|
]
|
|
)
|
|
X = np.empty((0, 2))
|
|
for i in range(8):
|
|
X = np.r_[X, means[i] + 0.8 * np.random.randn(n_features, 2)]
|
|
max_it = len(chunks)
|
|
it = 0
|
|
for chunk in chunks:
|
|
it += 1
|
|
print("==============================")
|
|
print("Iteration %03d of %03d" % (it, max_it))
|
|
print("==============================")
|
|
print()
|
|
|
|
print("Fast K-Means")
|
|
tstart = time()
|
|
mbkmeans = MiniBatchKMeans(init="k-means++", n_clusters=8, batch_size=chunk)
|
|
|
|
mbkmeans.fit(X)
|
|
delta = time() - tstart
|
|
print("Speed: %0.3fs" % delta)
|
|
print("Inertia: %0.3fs" % mbkmeans.inertia_)
|
|
print()
|
|
|
|
results["MiniBatchKMeans Speed"].append(delta)
|
|
results["MiniBatchKMeans Quality"].append(mbkmeans.inertia_)
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from mpl_toolkits.mplot3d import axes3d # noqa register the 3d projection
|
|
import matplotlib.pyplot as plt
|
|
|
|
samples_range = np.linspace(50, 150, 5).astype(int)
|
|
features_range = np.linspace(150, 50000, 5).astype(int)
|
|
chunks = np.linspace(500, 10000, 15).astype(int)
|
|
|
|
results = compute_bench(samples_range, features_range)
|
|
results_2 = compute_bench_2(chunks)
|
|
|
|
max_time = max(
|
|
[max(i) for i in [t for (label, t) in results.items() if "speed" in label]]
|
|
)
|
|
max_inertia = max(
|
|
[max(i) for i in [t for (label, t) in results.items() if "speed" not in label]]
|
|
)
|
|
|
|
fig = plt.figure("scikit-learn K-Means benchmark results")
|
|
for c, (label, timings) in zip("brcy", sorted(results.items())):
|
|
if "speed" in label:
|
|
ax = fig.add_subplot(2, 2, 1, projection="3d")
|
|
ax.set_zlim3d(0.0, max_time * 1.1)
|
|
else:
|
|
ax = fig.add_subplot(2, 2, 2, projection="3d")
|
|
ax.set_zlim3d(0.0, max_inertia * 1.1)
|
|
|
|
X, Y = np.meshgrid(samples_range, features_range)
|
|
Z = np.asarray(timings).reshape(samples_range.shape[0], features_range.shape[0])
|
|
ax.plot_surface(X, Y, Z.T, cstride=1, rstride=1, color=c, alpha=0.5)
|
|
ax.set_xlabel("n_samples")
|
|
ax.set_ylabel("n_features")
|
|
|
|
i = 0
|
|
for c, (label, timings) in zip("br", sorted(results_2.items())):
|
|
i += 1
|
|
ax = fig.add_subplot(2, 2, i + 2)
|
|
y = np.asarray(timings)
|
|
ax.plot(chunks, y, color=c, alpha=0.8)
|
|
ax.set_xlabel("Chunks")
|
|
ax.set_ylabel(label)
|
|
|
|
plt.show()
|