Minibatch convergence done, metric support left

Author: PyDataBlog

src/mini_batch.jl

@@ -13,7 +13,7 @@ MiniBatch() = MiniBatch(100)
 function kmeans!(alg::MiniBatch, X, k;
                  weights = nothing, metric = Euclidean(), n_threads = Threads.nthreads(),
                  k_init = "k-means++", init = nothing, max_iters = 300,
-                 tol = 0, max_no_improvement = 10, verbose = false, rng = Random.GLOBAL_RNG)
+                 tol = eltype(X)(1e-6), max_no_improvement = 10, verbose = false, rng = Random.GLOBAL_RNG)
 
     # Get the type and dimensions of design matrix, X
     T = eltype(X)
@@ -27,15 +27,16 @@ function kmeans!(alg::MiniBatch, X, k;
 
     # Initialize nearest centers
     labels = Vector{Int}(undef, alg.b)
+    final_labels = Vector{Int}(undef, ncol)
 
     converged = false
     niters = 0
+    counter = 0
     J_previous = zero(T)
     J = zero(T)
 
     # TODO: Main Steps. Batch update centroids until convergence
     while niters <= max_iters
-        counter = 0
 
         # b examples picked randomly from X (Step 5 in paper)
         batch_rand_idx = isnothing(weights) ? rand(rng, 1:ncol, alg.b) : wsample(rng, 1:ncol, weights, alg.b)
@@ -79,12 +80,27 @@ function kmeans!(alg::MiniBatch, X, k;
         end
 
         # TODO: Check for early stopping convergence
-        if (niters > 1) & abs(J - J_previous)
+        if (niters > 1) & abs((J - J_previous) < (tol * J))
             counter += 1
 
             # Declare convergence if max_no_improvement criterion is met
             if counter >= max_no_improvement
                 converged = true
+                # TODO: Compute label assignment for the complete dataset
+                @inbounds for i in axes(X, 2)
+                    min_dist = distance(metric, X, centroids, i, 1)
+                    label = 1
+
+                    for j in 2:size(centroids, 2)
+                        dist = distance(metric, X, centroids, i, j)
+                        label = dist < min_dist ? j : label
+                        min_dist = dist < min_dist ? dist : min_dist
+                    end
+
+                    final_labels[i] = label
+                end
+                # TODO: Compute totalcost for the complete dataset
+                J = sum_of_squares(X, final_labels, centroids)  # just a placeholder for now
                 break
             end
 
@@ -94,7 +110,7 @@ function kmeans!(alg::MiniBatch, X, k;
         niters += 1
     end
 
-    return centroids, niters, converged, labels, J  # TODO: push learned artifacts to KmeansResult
+    return centroids, niters, converged, final_labels, J  # TODO: push learned artifacts to KmeansResult
     #return KmeansResult(centroids, containers.labels, T[], Int[], T[], totalcost, niters, converged)
 end
 

test/test90_minibatch.jl

@@ -13,30 +13,30 @@ end
 
 
 @testset "MiniBatch convergence" begin
-    X = [1 1 1 4 4 4 4 0 2 3 5 1; 2 4 0 2 0 4 5 1 2 2 5 -1.]
-
     rng = StableRNG(2020)
-    baseline = kmeans(Lloyd(), X, 2, rng = rng)
+    X = rand(rng, 3, 100)
 
-    rng = StableRNG(2020)
-    res = kmeans(MiniBatch(6), X, 2, rng = rng)
+    baseline = [kmeans(Lloyd(), X, 2).totalcost for i in 1:1_000] |> mean |> round
+    # TODO: Switch to kmeans after full implementation
+    res = [ParallelKMeans.kmeans!(MiniBatch(50), X, 2)[end] for i in 1:1_000] |> mean |> round
 
-    @test baseline.totalcost ≈ res.totalcost
+    @test baseline == res
 end
 
 
 @testset "MiniBatch metric support" begin
-    X = [1 1 1 4 4 4 4 0 2 3 5 1; 2 4 0 2 0 4 5 1 2 2 5 -1.]
     rng = StableRNG(2020)
-    rng_orig = deepcopy(rng)
-
-    baseline = kmeans(Lloyd(), X, 2, tol = 1e-16, metric=Cityblock(), rng = rng)
+    X = rand(rng, 3, 100)
 
-    rng = deepcopy(rng_orig)
-    res = kmeans(MiniBatch(6), X, 2; tol = 1e-16, metric=Cityblock(), rng = rng)
+    baseline = [kmeans(Lloyd(), X, 2;
+                       tol=1e-6, metric=Cityblock(),
+                       max_iters=500).totalcost for i in 1:1000] |> mean |> floor
+    # TODO: Switch to kmeans after full implementation
+    res = [ParallelKMeans.kmeans!(MiniBatch(), X, 2;
+                                  metric=Cityblock(), tol=1e-6,
+                                  max_iters=500)[end] for i in 1:1000] |> mean |> floor
 
-    @test res.totalcost ≈ baseline.totalcost
-    @test res.converged == baseline.converged
+    @test baseline == res
 end