MLJ test

GPU launcher fixes
Evovest · Oct 26, 2023 · 57ac5dc · 57ac5dc
1 parent 9e5d4d6
commit 57ac5dc
Show file tree

Hide file tree

Showing 15 changed files with 21 additions and 6 deletions.
diff --git a/ext/EvoTreesCUDAExt/fit-utils.jl b/ext/EvoTreesCUDAExt/fit-utils.jl
@@ -30,8 +30,9 @@ function update_hist_gpu!(h, h∇_cpu, h∇, ∇, x_bin, is, js, jsc)
     max_blocks = config.blocks
     k = size(h∇, 1)
     ty = max(1, min(length(js), fld(max_threads, k)))
-    tx = max(1, min(length(is), fld(max_threads, k * ty)))
+    tx = min(64, max(1, min(length(is), fld(max_threads, k * ty))))
     threads = (k, ty, tx)
+    max_blocks = min(65535, max_blocks * fld(max_threads, prod(threads)))
     by = cld(length(js), ty)
     bx = min(cld(max_blocks, by), cld(length(is), tx))
     blocks = (1, by, bx)

diff --git a/ext/EvoTreesCUDAExt/predict.jl b/ext/EvoTreesCUDAExt/predict.jl
@@ -149,17 +149,19 @@ function EvoTrees.predict!(
     pred .= max.(T(-15), pred .- maximum(pred, dims=1))
 end
 
-# prediction from single tree - assign each observation to its final leaf
+# prediction for EvoTree model
 function predict(
     m::EvoTree{L,K},
     data,
     ::Type{<:EvoTrees.GPU};
     ntree_limit=length(m.trees)) where {L,K}
 
-    pred = CUDA.zeros(K, size(data, 1))
+    Tables.istable(data) ? data = Tables.columntable(data) : nothing
     ntrees = length(m.trees)
     ntree_limit > ntrees && error("ntree_limit is larger than number of trees $ntrees.")
     x_bin = CuArray(EvoTrees.binarize(data; fnames=m.info[:fnames], edges=m.info[:edges]))
+    nobs = size(x_bin, 1)
+    pred = CUDA.zeros(K, nobs)
     feattypes = CuArray(m.info[:feattypes])
     for i = 1:ntree_limit
         EvoTrees.predict!(pred, m.trees[i], x_bin, feattypes)

diff --git a/figures/gaussian-sinus-binary.png b/figures/gaussian-sinus-binary.png
diff --git a/figures/logistic-sinus-binary.png b/figures/logistic-sinus-binary.png
diff --git a/figures/quantiles-sinus-binary.png b/figures/quantiles-sinus-binary.png
diff --git a/figures/regression-sinus-binary-gpu.png b/figures/regression-sinus-binary-gpu.png
diff --git a/figures/regression-sinus-binary.png b/figures/regression-sinus-binary.png
diff --git a/figures/regression-sinus-oblivious-gpu.png b/figures/regression-sinus-oblivious-gpu.png
diff --git a/figures/regression-sinus-oblivious.png b/figures/regression-sinus-oblivious.png
diff --git a/figures/regression-sinus2-binary.png b/figures/regression-sinus2-binary.png
diff --git a/figures/regression-sinus2-oblivious.png b/figures/regression-sinus2-oblivious.png
diff --git a/src/MLJ.jl b/src/MLJ.jl
@@ -42,7 +42,7 @@ function MMI.update(
 end
 
 function predict(::EvoTreeRegressor, fitresult, A)
-  pred = vec(predict(fitresult, A))
+  pred = predict(fitresult, A)
   return pred
 end
 
@@ -52,7 +52,7 @@ function predict(::EvoTreeClassifier, fitresult, A)
 end
 
 function predict(::EvoTreeCount, fitresult, A)
-  λs = vec(predict(fitresult, A))
+  λs = predict(fitresult, A)
   return [Distributions.Poisson(λ) for λ ∈ λs]
 end
 

diff --git a/src/predict.jl b/src/predict.jl
@@ -86,10 +86,11 @@ function predict(
     ::Type{<:Device}=CPU;
     ntree_limit=length(m.trees)) where {L,K}
 
+    Tables.istable(data) ? data = Tables.columntable(data) : nothing
     ntrees = length(m.trees)
     ntree_limit > ntrees && error("ntree_limit is larger than number of trees $ntrees.")
     x_bin = binarize(data; fnames=m.info[:fnames], edges=m.info[:edges])
-    nobs = Tables.istable(data) ? length(Tables.getcolumn(data, 1)) : size(data, 1)
+    nobs = size(x_bin, 1)
     pred = zeros(Float32, K, nobs)
     for i = 1:ntree_limit
         predict!(pred, m.trees[i], x_bin, m.info[:feattypes])

diff --git a/test/MLJ.jl b/test/MLJ.jl
@@ -377,3 +377,13 @@ mach.model.nrounds += 10
 fit!(mach, rows=train, verbosity=1)
 
 report(mach)
+
+@testset "MLJ - rowtables - EvoTreeRegressor" begin
+    X, y = make_regression(1000, 5)
+    X = Tables.rowtable(X)
+    booster = EvoTreeRegressor()
+    # smoke tests:
+    mach = machine(booster, X, y) |> fit!
+    fit!(mach)
+    predict(mach, X)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -2,6 +2,7 @@ using Statistics
 using EvoTrees
 using EvoTrees: predict
 using CategoricalArrays
+using Tables
 using Random
 using Test