From 29b370562282ce9636638315bc69072aff76532a Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Tue, 20 Jan 2026 22:10:25 -0500 Subject: [PATCH 001/120] up --- benchmarks/benchmark_mse.jl | 332 ++++------------------------------ benchmarks/titanic-logloss.jl | 31 ++-- 2 files changed, 51 insertions(+), 312 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 25d9978..882a4e8 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -1,18 +1,10 @@ -# using Flux: params, gpu -using Revise - -using Tullio -using StatsBase using Statistics: mean, std using CUDA -using NeuroTrees +using NeuroTabModels using BenchmarkTools using Random: seed! -using ChainRulesCore -import ChainRulesCore: rrule - Threads.nthreads() seed!(123) @@ -21,294 +13,42 @@ num_feat = Int(100) @info "testing with: $nobs observations | $num_feat features." X = rand(Float32, nobs, num_feat) Y = randn(Float32, size(X, 1)) - -config = NeuroTreeRegressor( - device = :gpu, - loss = :mse, - nrounds = 1, - actA = :softmax, - scaler = false, - outsize = 1, - depth = 5, - num_trees = 64, - masks = nothing, - batchsize = 2048, - shuffle = true, - rng = 123, - opt = Dict("type" => "nadam", "lr" => 1e-3), +dtrain = DataFrame(X, :auto) +feature_names = names(dtrain) +dtrain.y = Y +target_name = "y" + +arch = NeuroTabModels.NeuroTreeConfig(; + actA=:identity, + init_scale=1.0, + depth=4, + ntrees=32, + stack_size=1, + hidden_size=1, ) - -CUDA.@time m, cache = NeuroTrees.init(config, x_train = X, y_train = Y); -CUDA.@time NeuroTreeModels.fit!(m, cache); -CUDA.@time NeuroTreeModels.fit(config, x_train = X, y_train = Y); -@time NeuroTrees.fit(config, x_train = X, y_train = Y); -CUDA.@time NeuroTreeModels.fit( - config, - x_train = X, - y_train = Y, - x_eval = X, - y_eval = Y, - metric = :mse, -); - -####### -# mse -####### -# scaler true - 1e6 cpu: -# scaler true - 1e6 gpu: -# No callback -# scaler true: 1e6 cpu: -# scaler true: 1e6 gpu: -_device = config.device == "cpu" ? NeuroTrees.cpu : NeuroTrees.gpu -dinfer = NeuroTreeModels.DataLoader( - Matrix{Float32}(X') |> _device, - batchsize = config.batchsize, - shuffle = false, - partial = true, +# arch = NeuroTabModels.MLPConfig(; +# act=:relu, +# stack_size=1, +# hidden_size=64, +# ) + +learner = NeuroTabRegressor( + arch; + loss=:mse, + nrounds=20, + early_stopping_rounds=2, + lr=1e-2, + device=:gpu ) -@time pred = NeuroTreeModels.infer(m, dinfer); -@btime pred = NeuroTreeModels.infer($m, $dinfer); -####### -# mse -####### -# scaler true - 1e6 cpu: -# scaler true - 1e6 gpu: - -# No BatchNorm -# scaler true: 1e6 cpu: 1.119 s (366929 allocations: 512.58 MiB) -# scaler true: 1e6 gpu: - - - -#### forward speed test -x = rand(Float32, num_feat, config.batchsize) |> gpu; -y = rand(Float32, config.batchsize) |> gpu; - -# cpu: -# gpu: 12.371 ms (495 allocations: 28.27 KiB) -@time m(x); -@btime CUDA.@sync m($x); -# CUDA.@time m(x); - -# cpu: -# gpu: 1.102 ms (350 allocations: 19.12 KiB) -@time nw = NeuroTreeModels.node_weights(m, copy(x)); -@btime NeuroTreeModels.node_weights($m, $copy(x)); -@btime CUDA.@sync NeuroTreeModels.node_weights($m, $x); - -# cpu: -@time _, lw = NeuroTreeModels.leaf_weights!(nw); -# gpu pre: 10.282 ms (77 allocations: 4.70 KiB) -# gpu post: 6.574 ms (78 allocations: 4.81 KiB) -@btime CUDA.@sync NeuroTreeModels.leaf_weights!($nw); - -# cpu: -# gpu: 899.400 μs (130 allocations: 9.11 KiB) -@time pred = NeuroTreeModels.dot_prod_tullio!(lw, m.p); -@btime CUDA.@sync NeuroTreeModels.dot_prod_tullio!($lw, $m.p); - - -############################################ -# node_weights breakdown -############################################ -# cpu: 125.000 μs (26 allocations: 190.42 KiB) -# gpu: 55.500 μs (128 allocations: 7.38 KiB) -@time fw = m.actA(m.w) .* m.mask; -@btime m.actA(m.w) .* m.mask; -@btime CUDA.@sync m.actA(m.w) .* m.mask; - -# cpu: 138.100 μs (274 allocations: 15.30 KiB) -# gpu: 83.100 μs (44 allocations: 2.08 KiB) -@time feat_proj = NeuroTreeModels.feat_proj!(fw, x); -@btime NeuroTreeModels.feat_proj!($fw, $x); -@btime CUDA.@sync NeuroTreeModels.feat_proj!($fw, $x); - -# cpu: 376.214 ns (5 allocations: 320 bytes) -# gpu: 940.625 ns (3 allocations: 160 bytes) -@time feat_proj = reshape(feat_proj, size(m.b, 1), size(m.b, 2), :); -@btime reshape($feat_proj, size(m.b, 1), size(m.b, 2), :); -@btime CUDA.@sync reshape($feat_proj, size(m.b, 1), size(m.b, 2), :); - -# cpu: 52.900 μs (245 allocations: 18.14 KiB) -# gpu: 39.200 μs (44 allocations: 4.80 KiB) -@time nw = NeuroTreeModels.nw_scale!(feat_proj, m.s, m.b); -@btime NeuroTreeModels.nw_scale!($feat_proj, $m.s, $m.b); -@btime CUDA.@sync NeuroTreeModels.nw_scale!($feat_proj, $m.s, $m.b); - -# cpu: 275.200 μs (81 allocations: 8.25 KiB) -# gpu: 25.800 μs (4 allocations: 160 bytes) -@time NeuroTreeModels.sigmoid_act!(nw); -@btime NeuroTreeModels.sigmoid_act!($nw); -@btime CUDA.@sync NeuroTreeModels.sigmoid_act!($nw); - -########################## -# grads -########################## -x = rand(Float32, num_feat, config.batchsize) |> gpu; -y = rand(Float32, config.batchsize) |> gpu; -w = ones(Float32, config.batchsize) |> gpu; -θ = NeuroTrees.params(m) - -function forward_test_1(m, x) - for i = 1:488 - vec(m(x)) - end - return nothing -end - -function grad_test1(m, x, y, w, θ) - for i = 1:488 - gs = NeuroTreeModels.gradient(θ) do - NeuroTreeModels.mse(m, x, y, w) - end - end - return nothing -end - -# 1.486592 seconds (257.94 k CPU allocations: 14.050 MiB, 0.83% gc time) (5.86 k GPU allocations: 5.948 GiB, 3.63% memmgmt time) -@time forward_test_1(m, x) -@btime forward_test_1($m, $x) -CUDA.@time CUDA.@sync forward_test_1(m, x) -# 3.587317 seconds (995.79 k CPU allocations: 63.791 MiB, 0.77% gc time) (21.47 k GPU allocations: 11.557 GiB, 4.12% memmgmt time) -@time grad_test1(m, x, y, w, θ) -@btime grad_test1($m, $x, $y, $w, $θ) -CUDA.@time CUDA.@sync grad_test1(m, x, y, w, θ) - -function model_2(x) - nw = NeuroTreeModels.node_weights(m, x) - lw = NeuroTreeModels.leaf_weights!(m.cnw, nw) - pred = NeuroTreeModels.dot_prod_tullio!(lw, m.p) ./ size(m.p, 3) - return pred -end - -function test_forward_2(x) - for i = 1:488 - model_2(x) - end - return nothing -end -function test_loss_2(x, y) - for i = 1:488 - NeuroTreeModels.mse(vec(model_2(x)), y) - end - return nothing -end - -function grad_test2(x, y, θ) - for i = 1:488 - gs = NeuroTreeModels.gradient(θ) do - NeuroTreeModels.mse(vec(model_2(x)), y) - end - end - return nothing -end - -# 2.894989 seconds (213.37 k allocations: 2.785 GiB, 7.06% gc time, 0.23% compilation time) -@time test_forward_2(x) -# @btime test_forward_2($x) - -# 2.938949 seconds (198.51 k allocations: 2.788 GiB, 6.84% gc time) -@time test_loss_2(x, y) -# @btime test_loss_2($x, $y) - -# 7.582608 seconds (1.27 M allocations: 6.116 GiB, 5.35% gc time, 6.91% compilation time) -@time grad_test2(x, y, θ) -# @btime grad_test2($x, $y, $θ) - - - -####################### -# post node_neights -####################### -function model_3(nw) - lw = NeuroTreeModels.leaf_weights!(m.cnw, nw) - pred = NeuroTreeModels.dot_prod_tullio!(lw, m.p) ./ size(m.p, 3) - return pred -end - -function test_forward_3(nw) - for i = 1:488 - model_3(nw) - end - return nothing -end -function test_loss_3(nw, y) - for i = 1:488 - NeuroTreeModels.mse(vec(model_3(nw)), y) - end - return nothing -end - -function grad_test3(nw, y, θ) - for i = 1:488 - gs = NeuroTreeModels.gradient(θ) do - NeuroTreeModels.mse(vec(model_3(nw)), y) - end - end - return nothing -end - -# 198.227 ms (56696 allocations: 8.65 MiB) -@time test_forward_3(nw) -@btime test_forward_3($nw) - -# 200.415 ms (58743 allocations: 12.57 MiB) -@time test_loss_3(nw, y) -@btime test_loss_3($nw, $y) - -# 2.019 s (180016 allocations: 972.06 MiB) -@time grad_test3(nw, y, θ) -@btime grad_test3($nw, $y, $θ) - - -####################### -# post lw -####################### -function model_4(lw) - pred = NeuroTreeModels.dot_prod_tullio!(lw, m.p) ./ size(m.p, 3) - return pred -end - -function test_forward_4(lw) - for i = 1:488 - model_3(lw) - end - return nothing -end -function test_loss_4(lw, y) - for i = 1:488 - NeuroTreeModels.mse(vec(model_3(lw)), y) - end - return nothing -end - -function grad_test4(lw, y, θ) - for i = 1:488 - gs = NeuroTreeModels.gradient(θ) do - NeuroTreeModels.mse(vec(model_4(lw)), y) - end - end - return nothing -end - -# 57.251 ms (14156 allocations: 4.60 MiB) -@time test_forward_4(lw) -@btime test_forward_4($lw) - -# 61.229 ms (16107 allocations: 8.52 MiB) -@time test_loss_4(lw, y) -@btime test_loss_4($lw, $y) - -# 548.048 ms (103333 allocations: 49.62 MiB) -@time grad_test4(lw, y, θ) -@btime grad_test4($lw, $y, $θ) +# nrounds=20: 32 sec +@time m = NeuroTabModels.fit( + learner, + dtrain; + target_name, + feature_names, + print_every_n=10, +) -function data_loop(data) - for d in dtrain - end - return nothing -end -# 325.788 ms (4409 allocations: 408.22 MiB) -@time data_loop(dtrain) -@btime data_loop($dtrain) +# @time p_train = m(dtrain; device=:gpu) +@time p_train = m(dtrain) diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 6b25328..21c6531 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -5,7 +5,6 @@ using StatsBase: median using CategoricalArrays using Random using CategoricalArrays -using EvoCore.IOTools using OrderedCollections using NeuroTabModels @@ -53,23 +52,23 @@ learner = NeuroTabRegressor( nrounds=400, early_stopping_rounds=2, lr=1e-2, + device=:gpu ) -learner = NeuroTabRegressor(; - arch_name="NeuroTreeConfig", - arch_config=Dict( - :actA => :identity, - :init_scale => 1.0, - :depth => 4, - :ntrees => 32, - :stack_size => 1, - :hidden_size => 1), - loss=:logloss, - nrounds=400, - early_stopping_rounds=2, - lr=1e-2, -) - +# learner = NeuroTabRegressor(; +# arch_name="NeuroTreeConfig", +# arch_config=Dict( +# :actA => :identity, +# :init_scale => 1.0, +# :depth => 4, +# :ntrees => 32, +# :stack_size => 1, +# :hidden_size => 1), +# loss=:logloss, +# nrounds=400, +# early_stopping_rounds=2, +# lr=1e-2, +# ) m = NeuroTabModels.fit( learner, From 4ab436fe172df47d7066bbef14c7292c65bd3dc6 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Wed, 21 Jan 2026 00:15:50 -0500 Subject: [PATCH 002/120] up --- benchmarks/benchmark_mse.jl | 5 +- src/models/NeuroTree/model.jl | 172 ++++++++++++++--------------- src/models/NeuroTree/neurotrees.jl | 46 +++++++- 3 files changed, 130 insertions(+), 93 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 882a4e8..291f8ce 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -1,7 +1,5 @@ -using Statistics: mean, std -using CUDA - using NeuroTabModels +using DataFrames using BenchmarkTools using Random: seed! @@ -23,6 +21,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; init_scale=1.0, depth=4, ntrees=32, + proj_size=1, stack_size=1, hidden_size=1, ) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 6dcb36a..a21346e 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -1,71 +1,97 @@ -include("leaf_weights.jl") - -struct NeuroTree{W,B,P,F<:Function} - w::W - s::B - b::B +struct NeuroTree{DI,DP,M,P} + d_in::DI + d_proj::DP + mask::M p::P - actA::F - scaler::Bool -end -@layer NeuroTree - -function node_weights(m::NeuroTree, x) - # [N X T, F] * [F, B] => [N x T, B] - if m.scaler - nw = Flux.sigmoid_fast.(Flux.softplus.(m.s) .* (m.actA(m.w) * x .+ m.b)) - else - nw = Flux.sigmoid_fast.(m.actA(m.w) * x .+ m.b) - end - # [N x T, B] -> [N, T, B] - return reshape(nw, :, size(m.p, 3), size(x, 2)) end +@layer NeuroTree trainable = (d_in, d_proj, p) -function (m::NeuroTree{W,B,P,F})(x::W) where {W,B,P,F} - # [F, B] -> [N, T, B] - nw = node_weights(m, x) - # [N, T, B] -> [L, T, B] - (_, lw) = leaf_weights!(nw) - # [L, T, B], [P, L, T] -> [P, B] - pred = dot_prod_agg(lw, m.p) ./ size(m.p, 3) - return pred +function (m::NeuroTree)(x) + h = m.d_in(x) # [F,B] => [HNT,B] + h = reshape(h, size(m.d_proj.weight, 2), :) # [HNT,B] => [H,NTB] + nw = m.d_proj(h) # [H,NTB] => [1,NTB] + nw = reshape(nw, size(m.mask, 1), :) # [1,NTB] => [N,TB] + lw = softmax(m.mask' * nw) # [N,TB] => [L,TB] + lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] + p = m.p * lw ./ 16 # [LT,B] => [P,B] + return p end - -dot_prod_agg(lw, p) = dropdims(sum(reshape(lw, 1, size(lw)...) .* p, dims=(2, 3)), dims=(2, 3)) +# function (m::NeuroTree)(x) +# nw = m.d_in(x) # [F,B] => [NT,B] +# nw = reshape(nw, size(m.mask, 1), :) # [NT,B] => [N,TB] +# lw = softmax(m.mask' * nw) # [N,TB] => [L,TB] +# lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] +# p = m.p * lw ./ 16 # [LT,B] => [P,B] +# return p +# end """ - NeuroTree(; ins, outs, depth=4, ntrees=64, actA=identity, scaler=true, init_scale=1e-1) - NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; depth=4, ntrees=64, actA=identity, scaler=true, init_scale=1e-1) + NeuroTree(; ins, outs, depth=4, ntrees=64, actA=identity, init_scale=1e-1) + NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; depth=4, ntrees=64, actA=identity, init_scale=1e-1) Initialization of a NeuroTree. """ -function NeuroTree(; ins, outs, depth=4, ntrees=64, actA=identity, scaler=true, init_scale=1e-1) - nnodes = 2^depth - 1 - nleaves = 2^depth - nt = NeuroTree( - Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s - Float32.((rand(nnodes * ntrees) .- 0.5) ./ 4), # b - Float32.(randn(outs, nleaves, ntrees) .* init_scale), # p - actA, - scaler +function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) + mask = get_mask(Val(tree_type), depth) + nnodes = size(mask, 1) + nleaves = size(mask, 2) + + op = NeuroTree( + Dense(ins => proj_size * nnodes * ntrees, relu), # w + # Dense(ins => nnodes * ntrees), # w + Dense(proj_size => 1), # s + mask, + Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p ) - return nt + return op end -function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; depth=4, ntrees=64, actA=identity, scaler=true, init_scale=1e-1) - nnodes = 2^depth - 1 - nleaves = 2^depth - nt = NeuroTree( - Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s - Float32.((rand(nnodes * ntrees) .- 0.5) ./ 4), # b - Float32.(randn(outs, nleaves, ntrees) .* init_scale), # p - actA, - scaler +function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) + mask = get_mask(Val(tree_type), depth) + nnodes = size(mask, 1) + nleaves = size(mask, 2) + + op = NeuroTree( + Dense(ins => proj_size * nnodes * ntrees, relu), # w + # Dense(ins => nnodes * ntrees), # w + Dense(proj_size => 1), # s + mask, + Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p ) - return nt + return op end +function get_mask(::Val{:binary}, depth::Integer) + nodes = 2^depth - 1 + leaves = 2^depth + mask = zeros(Bool, nodes, leaves) + + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d) + stride = 2 * k + for b in 1:blocks + view(mask, 2^(d - 1) + b - 1, (b-1)*stride+1:(b-1)*stride+k,) .= true + end + end + return mask +end + +function get_mask(::Val{:oblivious}, depth::Integer) + leaves = 2^depth + mask = zeros(Bool, depth, leaves) + + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d) + stride = 2 * k + for b in 1:blocks + view(mask, d, (b-1)*stride+1:(b-1)*stride+k,) .= true + end + end + return mask +end + + """ StackTree A StackTree is made of a collection of NeuroTree. @@ -75,23 +101,23 @@ struct StackTree end @layer StackTree -function StackTree((ins, outs)::Pair{<:Integer,<:Integer}; depth=4, ntrees=64, stack_size=2, hidden_size=8, actA=identity, scaler=true, init_scale=1e-1) +function StackTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, stack_size=1, hidden_size=8, actA=identity, scaler=true, init_scale=1e-1) @assert stack_size == 1 || hidden_size >= outs trees = [] for i in 1:stack_size if i == 1 if i < stack_size - tree = NeuroTree(ins => hidden_size; depth, ntrees, actA, scaler, init_scale) + tree = NeuroTree(ins => hidden_size; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) push!(trees, tree) else - tree = NeuroTree(ins => outs; depth, ntrees, actA, scaler, init_scale) + tree = NeuroTree(ins => outs; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) push!(trees, tree) end elseif i < stack_size - tree = NeuroTree(hidden_size => hidden_size; depth, ntrees, actA, scaler, init_scale) + tree = NeuroTree(hidden_size => hidden_size; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) push!(trees, tree) else - tree = NeuroTree(hidden_size => outs; depth, ntrees, actA, scaler, init_scale) + tree = NeuroTree(hidden_size => outs; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) push!(trees, tree) end end @@ -111,31 +137,3 @@ function (m::StackTree)(x::AbstractMatrix) end return p end - - -function _identity_act(x) - return x ./ sum(abs.(x), dims=2) -end -function _tanh_act(x) - x = Flux.tanh_fast.(x) - return x ./ sum(abs.(x), dims=2) -end -function _hardtanh_act(x) - x = Flux.hardtanh.(x) - return x ./ sum(abs.(x), dims=2) -end - -""" - act_dict = Dict( - :identity => _identity_act, - :tanh => _tanh_act, - :hardtanh => _hardtanh_act, - ) - -Dictionary mapping features activation name to their function. -""" -const act_dict = Dict( - :identity => _identity_act, - :tanh => _tanh_act, - :hardtanh => _hardtanh_act, -) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index f1f5083..8f8a49e 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -7,7 +7,7 @@ using CUDA import Flux import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch +import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel using ChainRulesCore @@ -19,9 +19,11 @@ import ..Models: Architecture include("model.jl") struct NeuroTreeConfig <: Architecture + tree_type::Symbol actA::Symbol depth::Int ntrees::Int + proj_size::Int hidden_size::Int stack_size::Int scaler::Bool @@ -33,9 +35,11 @@ function NeuroTreeConfig(; kwargs...) # defaults arguments args = Dict{Symbol,Any}( - :actA => :tanh, + :tree_type => :binary, + :actA => :identity, :depth => 4, - :ntrees => 64, + :ntrees => 32, + :proj_size => 1, :hidden_size => 1, :stack_size => 1, :scaler => true, @@ -59,9 +63,11 @@ function NeuroTreeConfig(; kwargs...) end config = NeuroTreeConfig( + Symbol(args[:tree_type]), Symbol(args[:actA]), args[:depth], args[:ntrees], + args[:proj_size], args[:hidden_size], args[:stack_size], args[:scaler], @@ -81,16 +87,20 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) Parallel( vcat, StackTree(nfeats => outsize; + tree_type=config.tree_type, depth=config.depth, ntrees=config.ntrees, + proj_size=config.proj_size, stack_size=config.stack_size, hidden_size=config.hidden_size, actA=act_dict[config.actA], scaler=config.scaler, init_scale=config.init_scale), StackTree(nfeats => outsize; + tree_type=config.tree_type, depth=config.depth, ntrees=config.ntrees, + proj_size=config.proj_size, stack_size=config.stack_size, hidden_size=config.hidden_size, actA=act_dict[config.actA], @@ -102,8 +112,10 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) chain = Chain( BatchNorm(nfeats), StackTree(nfeats => outsize; + tree_type=config.tree_type, depth=config.depth, ntrees=config.ntrees, + proj_size=config.proj_size, stack_size=config.stack_size, hidden_size=config.hidden_size, actA=act_dict[config.actA], @@ -115,4 +127,32 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) end +function _identity_act(x) + return x ./ sum(abs.(x), dims=2) +end +function _tanh_act(x) + x = Flux.tanh_fast.(x) + return x ./ sum(abs.(x), dims=2) +end +function _hardtanh_act(x) + x = Flux.hardtanh.(x) + return x ./ sum(abs.(x), dims=2) +end + +""" + act_dict = Dict( + :identity => _identity_act, + :tanh => _tanh_act, + :hardtanh => _hardtanh_act, + ) + +Dictionary mapping features activation name to their function. +""" +const act_dict = Dict( + :identity => _identity_act, + :tanh => _tanh_act, + :hardtanh => _hardtanh_act, +) + + end \ No newline at end of file From 57e10fb166e8558dac21364f37c99a6d7d0d19fc Mon Sep 17 00:00:00 2001 From: Jeremie Desgagne-Bouchard Date: Wed, 21 Jan 2026 01:52:36 -0500 Subject: [PATCH 003/120] Update model.jl --- src/models/NeuroTree/model.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index a21346e..0ca97cc 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -13,7 +13,7 @@ function (m::NeuroTree)(x) nw = reshape(nw, size(m.mask, 1), :) # [1,NTB] => [N,TB] lw = softmax(m.mask' * nw) # [N,TB] => [L,TB] lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] - p = m.p * lw ./ 16 # [LT,B] => [P,B] + p = m.p * lw ./ size(m.mask, 2) # [LT,B] => [P,B] return p end # function (m::NeuroTree)(x) @@ -21,7 +21,7 @@ end # nw = reshape(nw, size(m.mask, 1), :) # [NT,B] => [N,TB] # lw = softmax(m.mask' * nw) # [N,TB] => [L,TB] # lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] -# p = m.p * lw ./ 16 # [LT,B] => [P,B] +# p = m.p * lw ./ size(m.mask, 2) # [LT,B] => [P,B] # return p # end From 5ea33c4d24e2936321183d96bd92de3f05adc502 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Fri, 23 Jan 2026 17:27:30 -0500 Subject: [PATCH 004/120] cleanup --- benchmarks/YEAR-regression.jl | 41 +++---- benchmarks/titanic-logloss.jl | 16 +-- experiments/tree-tensor.jl | 136 +++++++++++++++++++++++ src/Fit/fit.jl | 1 + src/models/NeuroTree/leaf_weights.jl | 154 --------------------------- src/models/NeuroTree/model.jl | 122 ++++++++++++++++----- src/models/NeuroTree/neurotrees.jl | 3 - 7 files changed, 262 insertions(+), 211 deletions(-) create mode 100644 experiments/tree-tensor.jl delete mode 100644 src/models/NeuroTree/leaf_weights.jl diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index b6d5d54..2977760 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -41,38 +41,39 @@ dtrain = df_tot[train_idx, :]; deval = df_tot[eval_idx, :]; dtest = df_tot[(end-51630+1):end, :]; -# arch = NeuroTabModels.NeuroTreeConfig(; -# actA=:identity, -# depth=4, -# ntrees=32, -# stack_size=1, -# hidden_size=1, -# init_scale=0.1, -# MLE_tree_split=false -# ) +arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + proj_size=4, + actA=:identity, + depth=4, + ntrees=32, + stack_size=1, + hidden_size=1, + init_scale=1.0, + MLE_tree_split=false +) # arch = NeuroTabModels.MLPConfig(; # act=:relu, # stack_size=1, # hidden_size=256, # ) -arch = NeuroTabModels.ResNetConfig(; - num_blocks=1, - hidden_size=128, - act=:relu, - dropout=0.5, - MLE_tree_split=false -) +# arch = NeuroTabModels.ResNetConfig(; +# num_blocks=1, +# hidden_size=128, +# act=:relu, +# dropout=0.5, +# MLE_tree_split=false +# ) -device = :gpu -# :mse :gaussian_mle :tweedie -loss = :mse +device = :cpu +loss = :mse # :mse :gaussian_mle :tweedie learner = NeuroTabRegressor( arch; loss, nrounds=200, early_stopping_rounds=2, - lr=3e-4, + lr=1e-3, batchsize=1024, device ) diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 21c6531..51f5f7d 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -33,12 +33,14 @@ target_name = "Survived" feature_names = setdiff(names(df), ["Survived"]) arch = NeuroTabModels.NeuroTreeConfig(; - actA=:identity, + tree_type=:binary, + proj_size=1, init_scale=1.0, depth=4, - ntrees=32, + ntrees=16, stack_size=1, hidden_size=1, + actA=:identity, ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, @@ -49,10 +51,10 @@ arch = NeuroTabModels.NeuroTreeConfig(; learner = NeuroTabRegressor( arch; loss=:logloss, - nrounds=400, + nrounds=200, early_stopping_rounds=2, - lr=1e-2, - device=:gpu + lr=3e-2, + device=:cpu ) # learner = NeuroTabRegressor(; @@ -70,14 +72,14 @@ learner = NeuroTabRegressor( # lr=1e-2, # ) -m = NeuroTabModels.fit( +@time m = NeuroTabModels.fit( learner, dtrain; deval, target_name, feature_names, print_every_n=10, -) +); p_train = m(dtrain) p_eval = m(deval) diff --git a/experiments/tree-tensor.jl b/experiments/tree-tensor.jl new file mode 100644 index 0000000..e099dd4 --- /dev/null +++ b/experiments/tree-tensor.jl @@ -0,0 +1,136 @@ + +using BenchmarkTools +using Random +using NeuroTabModels +using NeuroTabModels.Models.NeuroTrees +using CUDA + +using CairoMakie +# density(m.chain.layers[2].trees[1].b) +# density(m.chain.layers[2].trees[1].s) +# mean(m.chain.layers[2].trees[1].s) +# density(vec(m.chain.layers[2].trees[1].p)) +# density(vec(m.chain.layers[2].trees[1].w)) +# density(abs.(vec(m.chain.layers[2].trees[1].w))) +# mean(abs.(vec(m.chain.layers[2].trees[1].w)) .< 1e-1) + +############################### +# original mask +############################### +mask = NeuroTrees.get_mask(Val(:binary), 4) +mask = NeuroTrees.get_mask(Val(:oblivious), 4) + +fig = Figure(; size=(450, 450)) +ax = Axis(fig[1, 1]; + title="original mask", + xlabel="Nodes", + ylabel="Leaves", + aspect=DataAspect(), + xticks=collect(1:size(mask, 1)), + yticks=collect(1:size(mask, 2)), + xgridcolor=:lightgrey, + ygridcolor=:lightgrey +) +hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) +translate!(hm, 0, 0, -100) +fig + +############################### +# softplus mask +############################### +mask = NeuroTrees.get_softplus_mask(Val(:binary), 4) +mask = NeuroTrees.get_softplus_mask(Val(:oblivious), 4) + +fig = Figure(; size=(450, 450)) +ax = Axis(fig[1, 1]; + title="softplus mask", + xlabel="Leaves", + ylabel="Nodes", + aspect=DataAspect(), + xticks=collect(1:size(mask, 1)), + yticks=collect(1:size(mask, 2)), + xgridcolor=:lightgrey, + ygridcolor=:lightgrey +) +hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) +translate!(hm, 0, 0, -100) +fig + +############################### +# logits mask +############################### +mask = NeuroTrees.get_logits_mask(Val(:binary), 4) +mask = NeuroTrees.get_logits_mask(Val(:oblivious), 4) + +fig = Figure(; size=(450, 450)) +ax = Axis(fig[1, 1]; + title="logits mask", + xlabel="Leaves", + ylabel="Nodes", + aspect=DataAspect(), + xticks=collect(1:size(mask, 1)), + yticks=collect(1:size(mask, 2)), + xgridcolor=:lightgrey, + ygridcolor=:lightgrey +) +hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) +translate!(hm, 0, 0, -100) +fig + +function leaf_weights(logits, lmask, smask) + logits * lmask .+ softplus.(logits) * smask +end +############## +# test +############## +depth = 3 +tree_type = :binary +nw = tree_type == :oblivious ? zeros(depth) : zeros(2^depth - 1) +nw[2] = -1.0 +lmask = NeuroTrees.get_logits_mask(Val(tree_type), depth) +smask = NeuroTrees.get_softplus_mask(Val(tree_type), depth) +lw = exp.(lmask * nw .- smask * NeuroTrees.softplus.(nw)) # [N,TB] => [L,TB] +sum(lw) + + +######################## +# benchmarks +######################## +function matmul_version(nw, mask) + mask * nw +end + +function broadcast_version(nw, mask) + # reshape(nw, 1, :) is the same as nw' in row form for broadcasting + dropdims(sum(mask .* reshape(nw, 1, size(nw)...); dims=2); dims=2) +end + +Random.seed!(123) +tree_type = :oblivious +depth = 2 + +for (n_leaves, depth) in [(32, 5), (64, 6)] + println("\n=== n_leaves = $n_leaves depth = $depth (density ≈ 50%) ===") + + nw = randn(Float32, depth, 1024) + nw = nw |> CuArray + + lmask = NeuroTrees.get_logits_mask(Val(tree_type), depth) + # lmask = lmask |> CuArray + # lmask = Float32.(lmask) |> CuArray + lmask = BitMatrix(lmask) |> CuArray + + # Warmup + matmul_version(nw, lmask) + broadcast_version(nw, lmask) + + b_mat = @benchmark matmul_version($nw, $lmask) + b_bc = @benchmark broadcast_version($nw, $lmask) + + t_mat = median(b_mat.times) / 1e3 # μs + t_bc = median(b_bc.times) / 1e3 # μs + + println(" Matmul: $(round(t_mat; digits=2)) μs (± $(round(std(b_mat.times)/1e3; digits=2)) μs)") + println(" Broadcast: $(round(t_bc; digits=2)) μs (± $(round(std(b_bc.times)/1e3; digits=2)) μs)") + println(" Ratio bc / matmul = $(round(t_bc / t_mat; digits=3))×") +end diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 68d13b6..1c1b552 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -62,6 +62,7 @@ function init( end optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + # optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) diff --git a/src/models/NeuroTree/leaf_weights.jl b/src/models/NeuroTree/leaf_weights.jl deleted file mode 100644 index cbe588a..0000000 --- a/src/models/NeuroTree/leaf_weights.jl +++ /dev/null @@ -1,154 +0,0 @@ -""" - leaf_weights!(cw, nw) - -Compute the cumulative probability associated with each node, down to terminal leaves -""" -function leaf_weights!(nw) - cw = ones(eltype(nw), 2 * size(nw, 1) + 1, size(nw)[2:3]...) - @threads for batch in axes(nw, 3) - @inbounds for tree in axes(nw, 2) - for i = 2:2:size(cw, 1) - # child cumulative probability is obtained from the product of the parent cumulative prob (cw[parent]) and the parent probability (np[parent]) - cw[i, tree, batch] = cw[i>>1, tree, batch] * nw[i>>1, tree, batch] - cw[i+1, tree, batch] = cw[i>>1, tree, batch] * (1 - nw[i>>1, tree, batch]) - end - end - end - @views lw = cw[size(nw, 1)+1:size(cw, 1), :, 1:size(nw, 3)] - return (cw, lw) -end - -""" - leaf_weights!(cw::AnyCuArray, nw::AnyCuArray) - -Compute the cumulative probability associated with each node, down to terminal leaves. -""" -function leaf_weights!(nw::AnyCuArray) - cw = CUDA.ones(eltype(nw), 2 * size(nw, 1) + 1, size(nw)[2:3]...) - blocks = size(nw, 3) - threads = size(nw, 2) - @cuda threads = threads blocks = blocks leaf_weights!_kernel!(cw, nw) - CUDA.synchronize() - @views lw = cw[size(nw, 1)+1:size(cw, 1), :, 1:size(nw, 3)] - return (cw, lw) -end - -function leaf_weights!_kernel!(cw::CuDeviceArray, nw::CuDeviceArray) - - tree = threadIdx().x # one thread per tree - batch = blockIdx().x # one block per batch - - @inbounds for i in 2:2:size(cw, 1) - # child cumulative probability is obtained from the product of the parent cumulative prob (cw[parent]) and the parent probability (np[parent]) - parent = cw[i>>1, tree, batch] - child = nw[i>>1, tree, batch] - cw[i, tree, batch] = parent * child - cw[i+1, tree, batch] = parent * (1 - child) - end - return nothing -end - -""" - rrule(::leaf_weights!, cw, nw) - -Backpropagation rule for the leaf probability calculation function -""" -function rrule(::typeof(leaf_weights!), nw) - cw, lw = leaf_weights!(nw) - max_depth = floor(Int, log2(size(nw, 1) + 1)) - node_offset = size(nw, 1) # offset on the leaf row - leaf_weights!_pullback(ȳ) = - NoTangent(), Δ_leaf_weights!(unthunk(ȳ[2]), cw, nw, max_depth, node_offset) - return (cw, lw), leaf_weights!_pullback -end - -""" - Δ_leaf_weights!(Δnw, ȳ, cw, nw, max_depth, node_offset) - -Kernel launcher for backpropagation rule of the leaf probability calculation -""" -function Δ_leaf_weights!(ȳ, cw, nw, max_depth, node_offset) - Δnw = zeros(eltype(nw), size(nw)...) - Δ_leaf_weights_kernel!(Δnw, ȳ, cw, nw, max_depth, node_offset) - return Δnw -end -function Δ_leaf_weights!(ȳ, cw::AnyCuArray, nw::AnyCuArray, max_depth, node_offset) - Δnw = CUDA.zeros(eltype(nw), size(nw)...) - blocks = size(nw, 3) - threads = size(nw, 2) - @cuda threads = threads blocks = blocks Δ_leaf_weights_kernel!( - Δnw, - ȳ, - cw, - nw, - max_depth, - node_offset, - ) - CUDA.synchronize() - return Δnw -end - -""" - leaf_weights!_pullback_kernel!(nw̄, ȳ, cw, nw, max_depth, node_offset) - -Kernel for backpropagation rule of the leaf probability calculation -""" -function Δ_leaf_weights_kernel!(Δnw, ȳ, cw, nw, max_depth, node_offset) - - @threads for batch in axes(nw, 3) - @inbounds for tree in axes(nw, 2) - # loop on node weights - for i in axes(nw, 1) - depth = floor(Int, log2(i)) # current depth level - starting at 0 - step = 2^(max_depth - depth) # iteration length - leaf_offset = step * (i - 2^depth) # offset on the leaf row - # loop on node weight leaf dependencies - half positive + half negative - for j = (1+leaf_offset):(step÷2+leaf_offset) - k = j + node_offset # move from leaf position to full tree position - Δnw[i, tree, batch] += - ȳ[j, tree, batch] * cw[k, tree, batch] / - max(1.0f-8, nw[i, tree, batch]) - end - for j = (1+leaf_offset+step÷2):(step+leaf_offset) - k = j + node_offset - Δnw[i, tree, batch] -= - ȳ[j, tree, batch] * cw[k, tree, batch] / - max(1.0f-8, (1 - nw[i, tree, batch])) - end - end - end - end - return nothing -end - -function Δ_leaf_weights_kernel!( - Δnw::CuDeviceArray, - ȳ, - cw::CuDeviceArray, - nw::CuDeviceArray, - max_depth, - node_offset, -) - - tree = threadIdx().x # one thread per tree - batch = blockIdx().x # one block per batch - - @inbounds for i in axes(nw, 1) - depth = floor(Int, log2(i)) # current depth level - starting at 0 - step = 2^(max_depth - depth) # iteration length - leaf_offset = step * (i - 2^depth) # offset on the leaf row - child = nw[i, tree, batch] - # loop on node weight leaf dependencies - half positive + half negative - @inbounds for j in (1+leaf_offset:step÷2+leaf_offset) - k = j + node_offset # move from leaf position to full tree position - Δnw[i, tree, batch] += - ȳ[j, tree, batch] * cw[k, tree, batch] / max(1.0f-8, child) - end - @inbounds for j = (1+leaf_offset+step÷2):(step+leaf_offset) - k = j + node_offset - Δnw[i, tree, batch] -= - ȳ[j, tree, batch] * cw[k, tree, batch] / max(1.0f-8, (1 - child)) - end - end - return nothing -end diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 0ca97cc..7450ded 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -1,29 +1,39 @@ struct NeuroTree{DI,DP,M,P} + ntrees::Int d_in::DI d_proj::DP - mask::M + lmask::M + smask::M p::P end @layer NeuroTree trainable = (d_in, d_proj, p) -function (m::NeuroTree)(x) - h = m.d_in(x) # [F,B] => [HNT,B] - h = reshape(h, size(m.d_proj.weight, 2), :) # [HNT,B] => [H,NTB] - nw = m.d_proj(h) # [H,NTB] => [1,NTB] - nw = reshape(nw, size(m.mask, 1), :) # [1,NTB] => [N,TB] - lw = softmax(m.mask' * nw) # [N,TB] => [L,TB] - lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] - p = m.p * lw ./ size(m.mask, 2) # [LT,B] => [P,B] - return p -end +# function (m::NeuroTree)(x) +# h = m.d_in(x) # [F,B] => [HNT,B] +# h = reshape(h, size(m.d_proj.weight, 2), :) # [HNT,B] => [H,NTB] +# nw = m.d_proj(h) # [H,NTB] => [1,NTB] +# nw = reshape(nw, size(m.lmask, 2), :) # [1,NTB] => [N,TB] +# lw = exp.(m.lmask * nw .- m.smask * softplus.(nw)) # [N,TB] => [L,TB] +# lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] +# p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] +# return p +# end # function (m::NeuroTree)(x) # nw = m.d_in(x) # [F,B] => [NT,B] -# nw = reshape(nw, size(m.mask, 1), :) # [NT,B] => [N,TB] -# lw = softmax(m.mask' * nw) # [N,TB] => [L,TB] +# nw = reshape(nw, size(m.lmask, 2), :) # [NT,B] => [N,TB] +# lw = exp.(m.lmask * nw .- m.smask * softplus.(nw)) # [N,TB] => [L,TB] # lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] -# p = m.p * lw ./ size(m.mask, 2) # [LT,B] => [P,B] +# p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] # return p # end +function (m::NeuroTree)(x) + nw = relu.(m.d_in(x)) # [F,B] => [NT,B] + nw = reshape(nw, size(m.lmask, 2), :) # [NT,B] => [N,TB] + lw = exp.(m.lmask * nw .- m.smask * softplus.(nw)) # [N,TB] => [L,TB] + lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] + p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] + return p +end """ NeuroTree(; ins, outs, depth=4, ntrees=64, actA=identity, init_scale=1e-1) @@ -32,60 +42,118 @@ end Initialization of a NeuroTree. """ function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) - mask = get_mask(Val(tree_type), depth) - nnodes = size(mask, 1) - nleaves = size(mask, 2) + lmask = get_logits_mask(Val(tree_type), depth) + smask = get_softplus_mask(Val(tree_type), depth) + nnodes = size(lmask, 1) + nleaves = size(lmask, 2) op = NeuroTree( + ntrees, Dense(ins => proj_size * nnodes * ntrees, relu), # w # Dense(ins => nnodes * ntrees), # w Dense(proj_size => 1), # s - mask, + Float32.(lmask), + Float32.(smask), Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p ) return op end function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) - mask = get_mask(Val(tree_type), depth) - nnodes = size(mask, 1) - nleaves = size(mask, 2) + lmask = get_logits_mask(Val(tree_type), depth) + smask = get_softplus_mask(Val(tree_type), depth) + nleaves = size(lmask, 1) + nnodes = size(lmask, 2) op = NeuroTree( + ntrees, Dense(ins => proj_size * nnodes * ntrees, relu), # w # Dense(ins => nnodes * ntrees), # w Dense(proj_size => 1), # s - mask, + Float32.(lmask), + Float32.(smask), Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p ) return op end + +""" + get_logits_mask(::Val{:binary}, depth::Integer) +""" +function get_logits_mask(::Val{:binary}, depth::Integer) + nodes = 2^depth - 1 + leaves = 2^depth + mask = zeros(Bool, leaves, nodes) + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d) + stride = 2 * k + for b in 1:blocks + view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= true + end + end + return mask +end +function get_logits_mask(::Val{:oblivious}, depth::Integer) + leaves = 2^depth + mask = zeros(Bool, leaves, depth) + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d) + stride = 2 * k + for b in 1:blocks + view(mask, (b-1)*stride+1:(b-1)*stride+k, d) .= true + end + end + return mask +end + +""" + get_softplus_mask(::Val{:binary}, depth::Integer) +""" +function get_softplus_mask(::Val{:binary}, depth::Integer) + nodes = 2^depth - 1 + leaves = 2^depth + mask = zeros(Bool, leaves, nodes) + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d + 1) + stride = k + for b in 1:blocks + view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= true + end + end + return mask +end +function get_softplus_mask(::Val{:oblivious}, depth::Integer) + leaves = 2^depth + mask = ones(Bool, leaves, depth) + return mask +end + function get_mask(::Val{:binary}, depth::Integer) nodes = 2^depth - 1 leaves = 2^depth mask = zeros(Bool, nodes, leaves) - for d in 1:depth blocks = 2^(d - 1) k = 2^(depth - d) stride = 2 * k for b in 1:blocks - view(mask, 2^(d - 1) + b - 1, (b-1)*stride+1:(b-1)*stride+k,) .= true + view(mask, 2^(d - 1) + b - 1, (b-1)*stride+1:(b-1)*stride+k) .= true end end return mask end - function get_mask(::Val{:oblivious}, depth::Integer) leaves = 2^depth mask = zeros(Bool, depth, leaves) - for d in 1:depth blocks = 2^(d - 1) k = 2^(depth - d) stride = 2 * k for b in 1:blocks - view(mask, d, (b-1)*stride+1:(b-1)*stride+k,) .= true + view(mask, d, (b-1)*stride+1:(b-1)*stride+k) .= true end end return mask diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 8f8a49e..758e65b 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -10,9 +10,6 @@ import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel -using ChainRulesCore -import ChainRulesCore: rrule - import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture From 502468848a61f5a5ab5cc9d6a43fbfdb4e9b85cf Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 23 Jan 2026 19:20:25 -0500 Subject: [PATCH 005/120] up --- Project.toml | 2 +- benchmarks/benchmark_mse.jl | 15 +++++++++------ experiments/tree-tensor.jl | 20 ++++++++++---------- src/models/NeuroTree/model.jl | 10 +++++----- 4 files changed, 25 insertions(+), 22 deletions(-) diff --git a/Project.toml b/Project.toml index ecca32c..5065c25 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NeuroTabModels" uuid = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" authors = ["jeremie.db "] -version = "0.2.3" +version = "0.3.0" [deps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 291f8ce..f1ed656 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -17,11 +17,12 @@ dtrain.y = Y target_name = "y" arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + proj_size=1, actA=:identity, init_scale=1.0, depth=4, ntrees=32, - proj_size=1, stack_size=1, hidden_size=1, ) @@ -34,20 +35,22 @@ arch = NeuroTabModels.NeuroTreeConfig(; learner = NeuroTabRegressor( arch; loss=:mse, - nrounds=20, + nrounds=10, early_stopping_rounds=2, lr=1e-2, + batchsize=2048, device=:gpu ) -# nrounds=20: 32 sec +# desktop gpu: 13.476383 seconds (26.42 M allocations: 5.990 GiB, 9.44% gc time) +# 13.557744 seconds (26.40 M allocations: 5.989 GiB, 9.60% gc time) @time m = NeuroTabModels.fit( learner, dtrain; target_name, feature_names, print_every_n=10, -) +); -# @time p_train = m(dtrain; device=:gpu) -@time p_train = m(dtrain) +# desktop: 0.771839 seconds (369.20 k allocations: 1.522 GiB, 5.94% gc time) +@time p_train = m(dtrain; device=:gpu); diff --git a/experiments/tree-tensor.jl b/experiments/tree-tensor.jl index e099dd4..9cb7960 100644 --- a/experiments/tree-tensor.jl +++ b/experiments/tree-tensor.jl @@ -4,8 +4,8 @@ using Random using NeuroTabModels using NeuroTabModels.Models.NeuroTrees using CUDA +# using CairoMakie -using CairoMakie # density(m.chain.layers[2].trees[1].b) # density(m.chain.layers[2].trees[1].s) # mean(m.chain.layers[2].trees[1].s) @@ -106,19 +106,19 @@ function broadcast_version(nw, mask) end Random.seed!(123) -tree_type = :oblivious -depth = 2 - -for (n_leaves, depth) in [(32, 5), (64, 6)] - println("\n=== n_leaves = $n_leaves depth = $depth (density ≈ 50%) ===") +tree_type = :binary +# depth = 2 - nw = randn(Float32, depth, 1024) - nw = nw |> CuArray +for depth in [4, 5, 6] + println("\n=== tree_type = $tree_type | depth = $depth ===") lmask = NeuroTrees.get_logits_mask(Val(tree_type), depth) # lmask = lmask |> CuArray - # lmask = Float32.(lmask) |> CuArray - lmask = BitMatrix(lmask) |> CuArray + lmask = Float32.(lmask) # |> CuArray + # lmask = BitMatrix(lmask) |> CuArray + + nw = randn(Float32, size(lmask, 2), 1024) + # nw = nw |> CuArray # Warmup matmul_version(nw, lmask) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 7450ded..51f68eb 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -89,7 +89,7 @@ function get_logits_mask(::Val{:binary}, depth::Integer) k = 2^(depth - d) stride = 2 * k for b in 1:blocks - view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= true + view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= 1 end end return mask @@ -102,7 +102,7 @@ function get_logits_mask(::Val{:oblivious}, depth::Integer) k = 2^(depth - d) stride = 2 * k for b in 1:blocks - view(mask, (b-1)*stride+1:(b-1)*stride+k, d) .= true + view(mask, (b-1)*stride+1:(b-1)*stride+k, d) .= 1 end end return mask @@ -120,7 +120,7 @@ function get_softplus_mask(::Val{:binary}, depth::Integer) k = 2^(depth - d + 1) stride = k for b in 1:blocks - view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= true + view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= 1 end end return mask @@ -134,13 +134,13 @@ end function get_mask(::Val{:binary}, depth::Integer) nodes = 2^depth - 1 leaves = 2^depth - mask = zeros(Bool, nodes, leaves) + mask = zeros(Float32, nodes, leaves) for d in 1:depth blocks = 2^(d - 1) k = 2^(depth - d) stride = 2 * k for b in 1:blocks - view(mask, 2^(d - 1) + b - 1, (b-1)*stride+1:(b-1)*stride+k) .= true + view(mask, 2^(d - 1) + b - 1, (b-1)*stride+1:(b-1)*stride+k) .= 1 end end return mask From 31776f43b92da977dc710df978b73b6bc24ee4bf Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 23 Jan 2026 21:09:17 -0500 Subject: [PATCH 006/120] up --- benchmarks/YEAR-regression.jl | 4 +- benchmarks/benchmark_mse.jl | 1 + src/Fit/fit.jl | 3 +- src/models/NeuroTree/model.jl | 103 +++++++++++++++++----------------- 4 files changed, 57 insertions(+), 54 deletions(-) diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 2977760..c5864e6 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -1,7 +1,6 @@ using Random using CSV using DataFrames -using StatsBase using Statistics: mean, std using NeuroTabModels using AWS: AWSCredentials, AWSConfig, @service @@ -50,6 +49,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; stack_size=1, hidden_size=1, init_scale=1.0, + scaler=true, MLE_tree_split=false ) # arch = NeuroTabModels.MLPConfig(; @@ -65,7 +65,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # MLE_tree_split=false # ) -device = :cpu +device = :gpu loss = :mse # :mse :gaussian_mle :tweedie learner = NeuroTabRegressor( diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index f1ed656..a4ba492 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -25,6 +25,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; ntrees=32, stack_size=1, hidden_size=1, + scaler=false, ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 1c1b552..679e86e 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -61,8 +61,7 @@ function init( m = m |> gpu end - optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) - # optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 51f68eb..2d46be7 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -1,35 +1,29 @@ -struct NeuroTree{DI,DP,M,P} +struct NeuroTree{M,V,F} + w::M + b::V + s::V + p::M + ml::M + ms::M + actA::F + scaler::Bool ntrees::Int - d_in::DI - d_proj::DP - lmask::M - smask::M - p::P end -@layer NeuroTree trainable = (d_in, d_proj, p) +@layer NeuroTree trainable = (w, b, s, p) -# function (m::NeuroTree)(x) -# h = m.d_in(x) # [F,B] => [HNT,B] -# h = reshape(h, size(m.d_proj.weight, 2), :) # [HNT,B] => [H,NTB] -# nw = m.d_proj(h) # [H,NTB] => [1,NTB] -# nw = reshape(nw, size(m.lmask, 2), :) # [1,NTB] => [N,TB] -# lw = exp.(m.lmask * nw .- m.smask * softplus.(nw)) # [N,TB] => [L,TB] -# lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] -# p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] -# return p -# end -# function (m::NeuroTree)(x) -# nw = m.d_in(x) # [F,B] => [NT,B] -# nw = reshape(nw, size(m.lmask, 2), :) # [NT,B] => [N,TB] -# lw = exp.(m.lmask * nw .- m.smask * softplus.(nw)) # [N,TB] => [L,TB] -# lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] -# p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] -# return p -# end +# h = m.d_in(x) # [F,B] => [HNT,B] +# h = reshape(h, size(m.d_proj.weight, 2), :) # [HNT,B] => [H,NTB] +# nw = m.d_proj(h) # [H,NTB] => [1,NTB] +# nw = reshape(nw, size(m.lmask, 2), :) # [1,NTB] => [N,TB] function (m::NeuroTree)(x) - nw = relu.(m.d_in(x)) # [F,B] => [NT,B] - nw = reshape(nw, size(m.lmask, 2), :) # [NT,B] => [N,TB] - lw = exp.(m.lmask * nw .- m.smask * softplus.(nw)) # [N,TB] => [L,TB] + nw = m.w * x .+ m.b # [F,B] => [NT,B] + if m.scaler + nw = softplus.(m.s) .* (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] + else + nw = m.actA(m.w) * x .+ m.b # [F,B] => [NT,B] + end + nw = reshape(nw, size(m.ml, 2), :) # [NT,B] => [N,TB] + lw = exp.(m.ml * nw .- m.ms * softplus.(nw)) # [N,TB] => [L,TB] lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] return p @@ -41,37 +35,46 @@ end Initialization of a NeuroTree. """ -function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) - lmask = get_logits_mask(Val(tree_type), depth) - smask = get_softplus_mask(Val(tree_type), depth) - nnodes = size(lmask, 1) - nleaves = size(lmask, 2) +function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1f0) + ml = get_logits_mask(Val(tree_type), depth) + ms = get_softplus_mask(Val(tree_type), depth) + nleaves = size(ml, 1) + nnodes = size(ml, 2) op = NeuroTree( + # Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w + # Float32.((zeros(nnodes * ntrees) .- 0.5) ./ 4), # b + glorot_uniform(nnodes * ntrees, ins), # w + zeros(Float32, nnodes * ntrees), # b + Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s + glorot_uniform(outs, nleaves * ntrees), # p + # Float32.((rand(outs, nleaves * ntrees) .- 0.5) .* sqrt(12) .* init_scale), # p + # Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p + Float32.(ml), + Float32.(ms), + actA, + scaler, ntrees, - Dense(ins => proj_size * nnodes * ntrees, relu), # w - # Dense(ins => nnodes * ntrees), # w - Dense(proj_size => 1), # s - Float32.(lmask), - Float32.(smask), - Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p ) return op end -function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) - lmask = get_logits_mask(Val(tree_type), depth) - smask = get_softplus_mask(Val(tree_type), depth) - nleaves = size(lmask, 1) - nnodes = size(lmask, 2) +function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1f0) + ml = get_logits_mask(Val(tree_type), depth) + ms = get_softplus_mask(Val(tree_type), depth) + nleaves = size(ml, 1) + nnodes = size(ml, 2) op = NeuroTree( + Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w + Float32.((rand(nnodes * ntrees) .- 0.5) ./ 4), # b + Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s + Float32.((rand(outs, nleaves * ntrees) .- 0.5) .* sqrt(12) .* init_scale), # p + # Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p + Float32.(ml), + Float32.(ms), + actA, + scaler, ntrees, - Dense(ins => proj_size * nnodes * ntrees, relu), # w - # Dense(ins => nnodes * ntrees), # w - Dense(proj_size => 1), # s - Float32.(lmask), - Float32.(smask), - Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p ) return op end From 0859d4897fa2dffaaf78b3fb118b9f0d0a86675b Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 23 Jan 2026 22:27:01 -0500 Subject: [PATCH 007/120] up --- benchmarks/YEAR-regression.jl | 2 +- src/models/NeuroTree/model.jl | 21 ++++++++------------- src/models/NeuroTree/neurotrees.jl | 2 +- 3 files changed, 10 insertions(+), 15 deletions(-) diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index c5864e6..e8199ad 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -48,7 +48,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; ntrees=32, stack_size=1, hidden_size=1, - init_scale=1.0, + init_scale=0.0, scaler=true, MLE_tree_split=false ) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 2d46be7..a2e4ad8 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -18,9 +18,9 @@ end function (m::NeuroTree)(x) nw = m.w * x .+ m.b # [F,B] => [NT,B] if m.scaler - nw = softplus.(m.s) .* (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] + nw = softplus(m.s) .* relu.(m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] else - nw = m.actA(m.w) * x .+ m.b # [F,B] => [NT,B] + nw = relu.(m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] end nw = reshape(nw, size(m.ml, 2), :) # [NT,B] => [N,TB] lw = exp.(m.ml * nw .- m.ms * softplus.(nw)) # [N,TB] => [L,TB] @@ -35,21 +35,17 @@ end Initialization of a NeuroTree. """ -function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1f0) +function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) ml = get_logits_mask(Val(tree_type), depth) ms = get_softplus_mask(Val(tree_type), depth) nleaves = size(ml, 1) nnodes = size(ml, 2) op = NeuroTree( - # Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - # Float32.((zeros(nnodes * ntrees) .- 0.5) ./ 4), # b - glorot_uniform(nnodes * ntrees, ins), # w + Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w zeros(Float32, nnodes * ntrees), # b Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s - glorot_uniform(outs, nleaves * ntrees), # p - # Float32.((rand(outs, nleaves * ntrees) .- 0.5) .* sqrt(12) .* init_scale), # p - # Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p + Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p Float32.(ml), Float32.(ms), actA, @@ -58,7 +54,7 @@ function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size ) return op end -function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1f0) +function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) ml = get_logits_mask(Val(tree_type), depth) ms = get_softplus_mask(Val(tree_type), depth) nleaves = size(ml, 1) @@ -66,10 +62,9 @@ function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, de op = NeuroTree( Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - Float32.((rand(nnodes * ntrees) .- 0.5) ./ 4), # b + zeros(Float32, nnodes * ntrees), # b Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s - Float32.((rand(outs, nleaves * ntrees) .- 0.5) .* sqrt(12) .* init_scale), # p - # Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p + Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p Float32.(ml), Float32.(ms), actA, diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 758e65b..396e23b 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -7,7 +7,7 @@ using CUDA import Flux import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch +import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch, glorot_uniform import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel import ..Losses: get_loss_type, GaussianMLE From 97bd3dadc4e1e7b4e780aa6bd60d733e9f8f4e55 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Sun, 25 Jan 2026 23:31:44 -0500 Subject: [PATCH 008/120] up --- src/Fit/fit.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 679e86e..5de91b3 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -9,7 +9,7 @@ using ..Losses using ..Metrics import MLJModelInterface: fit -import CUDA +import CUDA, cuDNN import Optimisers import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta import Flux: trainmode!, gradient, cpu, gpu From 4bbdc7998fd0278121c20e3c569f72415ae7767f Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Tue, 27 Jan 2026 23:52:40 -0500 Subject: [PATCH 009/120] up --- src/Fit/fit.jl | 2 +- src/models/NeuroTree/model.jl | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 5de91b3..8ceffa8 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -61,7 +61,7 @@ function init( m = m |> gpu end - optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index a2e4ad8..9d03e00 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -18,9 +18,9 @@ end function (m::NeuroTree)(x) nw = m.w * x .+ m.b # [F,B] => [NT,B] if m.scaler - nw = softplus(m.s) .* relu.(m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] + nw = softplus(m.s) .* (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] else - nw = relu.(m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] + nw = (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] end nw = reshape(nw, size(m.ml, 2), :) # [NT,B] => [N,TB] lw = exp.(m.ml * nw .- m.ms * softplus.(nw)) # [N,TB] => [L,TB] @@ -62,7 +62,7 @@ function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, de op = NeuroTree( Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - zeros(Float32, nnodes * ntrees), # b + Float32.((rand(nnodes * ntrees) .- 0.5) ./ 4), # b Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p Float32.(ml), From e2a3dfc322f5d0c2dc423d3a12819abfbe8dc3fa Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 30 Jan 2026 19:16:41 +0530 Subject: [PATCH 010/120] Replace Zygote with Enzyme for gradient computation - Add Enzyme.jl dependency - Implement compute_grads() using Enzyme.autodiff with runtime activity mode - Add train/test mode switching to handle BatchNorm mutation issues - Refactor mlogloss to use direct indexing instead of onehotbatch - Configure Enzyme strictAliasing in module __init__ Signed-off-by: AdityaPandeyCN --- Project.toml | 5 +++++ src/Fit/fit.jl | 25 +++++++++++++++++++++++-- src/losses.jl | 31 ++++++++++++++++++++++++------- src/models/models.jl | 1 + 4 files changed, 53 insertions(+), 9 deletions(-) diff --git a/Project.toml b/Project.toml index 5065c25..8c5e4e4 100644 --- a/Project.toml +++ b/Project.toml @@ -8,8 +8,10 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" @@ -24,11 +26,14 @@ CUDA = "5" CategoricalArrays = "1" ChainRulesCore = "1" DataFrames = "1.3" +Enzyme = "0.13" Flux = "0.14, 0.15, 0.16" Functors = "0.5" +MLDatasets = "0.5.13" MLJModelInterface = "1.2.1" MLUtils = "0.4" Optimisers = "0.4" +PlotlyLight = "0.13.1" Random = "1" Statistics = "1" StatsBase = "0.34" diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 8ceffa8..512c086 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -10,9 +10,11 @@ using ..Metrics import MLJModelInterface: fit import CUDA, cuDNN +import Enzyme +import Enzyme: Duplicated, Active, Const, Reverse import Optimisers import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta -import Flux: trainmode!, gradient, cpu, gpu +import Flux: trainmode!, testmode!, cpu, gpu using DataFrames using CategoricalArrays @@ -20,6 +22,10 @@ using CategoricalArrays include("callback.jl") using .CallBacks +function __init__() + Enzyme.API.strictAliasing!(false) +end + function init( config::LearnerTypes, df::AbstractDataFrame; @@ -132,6 +138,7 @@ function fit( if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) + testmode!(m) cb(logger, 0, m) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] else @@ -143,6 +150,7 @@ function fit( fit_iter!(m, cache) iter = m.info[:nrounds] if !isnothing(logger) + testmode!(m) cb(logger, iter, m) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] @@ -156,6 +164,19 @@ function fit( # return m |> cpu end +function compute_grads(loss, model, d::Tuple) + trainmode!(model) + model(d[1]) + testmode!(model) + + dmodel = Enzyme.make_zero(model) + ad = Enzyme.set_runtime_activity(Reverse) + + Enzyme.autodiff(ad, Const(loss), Active, Duplicated(model, dmodel), Const.(d)...) + + return dmodel +end + function fit_iter!(m, cache) loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] GC.gc(true) @@ -163,7 +184,7 @@ function fit_iter!(m, cache) CUDA.reclaim() end for d in data - grads = gradient(model -> loss(model, d...), m)[1] + grads = compute_grads(loss, m, d) Optimisers.update!(opts, m, grads) end m.info[:nrounds] += 1 diff --git a/src/losses.jl b/src/losses.jl index bec273e..709d42a 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -69,20 +69,37 @@ function tweedie(m, x, y, w, offset) ) / sum(w) end +function _mlogloss_from_logprobs(p, y) + n = length(y) + acc = zero(eltype(p)) + @inbounds for i in 1:n + acc -= p[Int(y[i]), i] + end + return acc / n +end + +function _mlogloss_from_logprobs(p, y, w) + acc = zero(eltype(p)) + wsum = zero(eltype(w)) + @inbounds for i in 1:length(y) + wi = w[i] + wsum += wi + acc -= p[Int(y[i]), i] * wi + end + return acc / wsum +end + function mlogloss(m, x, y) p = logsoftmax(m(x); dims=1) - k = size(p, 1) - mean(-sum(onehotbatch(y, 1:k) .* p; dims=1)) + return _mlogloss_from_logprobs(p, y) end function mlogloss(m, x, y, w) p = logsoftmax(m(x); dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) + return _mlogloss_from_logprobs(p, y, w) end function mlogloss(m, x, y, w, offset) p = logsoftmax(m(x) .+ offset; dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) + return _mlogloss_from_logprobs(p, y, w) end gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}) where {T} = @@ -125,4 +142,4 @@ const _loss_type_dict = Dict( get_loss_type(loss::Symbol) = _loss_type_dict[loss] -end \ No newline at end of file +end diff --git a/src/models/models.jl b/src/models/models.jl index 41f3ddd..4e4ffa9 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -17,6 +17,7 @@ struct NeuroTabModel{L<:LossType,C<:Chain} chain::C info::Dict{Symbol,Any} end +@layer NeuroTabModel @functor NeuroTabModel (chain,) include("NeuroTree/neurotrees.jl") From e877c065e850c12cc7d0ec7339f944f975e2629c Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sat, 31 Jan 2026 19:02:44 +0530 Subject: [PATCH 011/120] revert back loss logic Signed-off-by: AdityaPandeyCN --- Project.toml | 5 +---- src/Fit/fit.jl | 43 ++++++++++++------------------------------- src/losses.jl | 31 +++++++------------------------ 3 files changed, 20 insertions(+), 59 deletions(-) diff --git a/Project.toml b/Project.toml index 8c5e4e4..a74159a 100644 --- a/Project.toml +++ b/Project.toml @@ -11,7 +11,6 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" -MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" @@ -29,11 +28,9 @@ DataFrames = "1.3" Enzyme = "0.13" Flux = "0.14, 0.15, 0.16" Functors = "0.5" -MLDatasets = "0.5.13" MLJModelInterface = "1.2.1" MLUtils = "0.4" Optimisers = "0.4" -PlotlyLight = "0.13.1" Random = "1" Statistics = "1" StatsBase = "0.34" @@ -48,4 +45,4 @@ MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test", "MLDatasets", "MLJTestInterface", "MLJBase"] +test = ["Test", "MLDatasets", "MLJTestInterface", "MLJBase"] \ No newline at end of file diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 512c086..8bf39b0 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -10,11 +10,12 @@ using ..Metrics import MLJModelInterface: fit import CUDA, cuDNN -import Enzyme -import Enzyme: Duplicated, Active, Const, Reverse +import Enzyme: Duplicated, Const +import Enzyme.API import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta -import Flux: trainmode!, testmode!, cpu, gpu +import Optimisers: OptimiserChain, WeightDecay, Adam +import Flux +import Flux: gradient, cpu, gpu using DataFrames using CategoricalArrays @@ -22,8 +23,11 @@ using CategoricalArrays include("callback.jl") using .CallBacks -function __init__() - Enzyme.API.strictAliasing!(false) +function compute_grads(loss, model, batch) + dup_model = Duplicated(model) + const_args = map(Const, batch) + grads = Flux.gradient((m, args...) -> loss(m, args...), dup_model, const_args...) + return grads[1] end function init( @@ -67,14 +71,13 @@ function init( m = m |> gpu end - optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) return m, cache end - """ function fit( config::NeuroTypes, @@ -91,16 +94,11 @@ end device=:cpu, gpuID=0, ) - Training function of NeuroTabModels' internal API. - # Arguments - - `config::LearnerTypes` - `dtrain`: Must be `<:AbstractDataFrame` - # Keyword arguments - - `feature_names`: Required kwarg, a `Vector{Symbol}` or `Vector{String}` of the feature names. - `target_name` Required kwarg, a `Symbol` or `String` indicating the name of the target variable. - `weight_name=nothing` @@ -109,6 +107,7 @@ Training function of NeuroTabModels' internal API. - `print_every_n=9999` - `verbosity=1` """ + function fit( config::LearnerTypes, dtrain; @@ -133,24 +132,20 @@ function fit( m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) - # initialize callback and logger if tracking eval data logger = nothing if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) - testmode!(m) cb(logger, 0, m) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] else (verbosity > 0) && @info "Init training" end - # for iter = 1:config.nrounds while m.info[:nrounds] < config.nrounds fit_iter!(m, cache) iter = m.info[:nrounds] if !isnothing(logger) - testmode!(m) cb(logger, iter, m) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] @@ -161,20 +156,6 @@ function fit( m.info[:logger] = logger return m - # return m |> cpu -end - -function compute_grads(loss, model, d::Tuple) - trainmode!(model) - model(d[1]) - testmode!(model) - - dmodel = Enzyme.make_zero(model) - ad = Enzyme.set_runtime_activity(Reverse) - - Enzyme.autodiff(ad, Const(loss), Active, Duplicated(model, dmodel), Const.(d)...) - - return dmodel end function fit_iter!(m, cache) diff --git a/src/losses.jl b/src/losses.jl index 709d42a..bec273e 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -69,37 +69,20 @@ function tweedie(m, x, y, w, offset) ) / sum(w) end -function _mlogloss_from_logprobs(p, y) - n = length(y) - acc = zero(eltype(p)) - @inbounds for i in 1:n - acc -= p[Int(y[i]), i] - end - return acc / n -end - -function _mlogloss_from_logprobs(p, y, w) - acc = zero(eltype(p)) - wsum = zero(eltype(w)) - @inbounds for i in 1:length(y) - wi = w[i] - wsum += wi - acc -= p[Int(y[i]), i] * wi - end - return acc / wsum -end - function mlogloss(m, x, y) p = logsoftmax(m(x); dims=1) - return _mlogloss_from_logprobs(p, y) + k = size(p, 1) + mean(-sum(onehotbatch(y, 1:k) .* p; dims=1)) end function mlogloss(m, x, y, w) p = logsoftmax(m(x); dims=1) - return _mlogloss_from_logprobs(p, y, w) + k = size(p, 1) + sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) end function mlogloss(m, x, y, w, offset) p = logsoftmax(m(x) .+ offset; dims=1) - return _mlogloss_from_logprobs(p, y, w) + k = size(p, 1) + sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) end gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}) where {T} = @@ -142,4 +125,4 @@ const _loss_type_dict = Dict( get_loss_type(loss::Symbol) = _loss_type_dict[loss] -end +end \ No newline at end of file From 2fe0010abe516b2107317d0e229c014379a4d86a Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 31 Jan 2026 18:31:35 -0500 Subject: [PATCH 012/120] up --- benchmarks/benchmark_mse.jl | 6 +++--- src/Fit/fit.jl | 41 +++++++++++++++++++++++-------------- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index a4ba492..6c4e14c 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -6,7 +6,7 @@ using Random: seed! Threads.nthreads() seed!(123) -nobs = Int(1e6) +nobs = Int(1e5) num_feat = Int(100) @info "testing with: $nobs observations | $num_feat features." X = rand(Float32, nobs, num_feat) @@ -36,11 +36,11 @@ arch = NeuroTabModels.NeuroTreeConfig(; learner = NeuroTabRegressor( arch; loss=:mse, - nrounds=10, + nrounds=1, early_stopping_rounds=2, lr=1e-2, batchsize=2048, - device=:gpu + device=:cpu ) # desktop gpu: 13.476383 seconds (26.42 M allocations: 5.990 GiB, 9.44% gc time) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 8bf39b0..02069f1 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -10,7 +10,8 @@ using ..Metrics import MLJModelInterface: fit import CUDA, cuDNN -import Enzyme: Duplicated, Const +import Enzyme +import Enzyme: Duplicated, Const, Reverse, set_runtime_activity import Enzyme.API import Optimisers import Optimisers: OptimiserChain, WeightDecay, Adam @@ -23,13 +24,6 @@ using CategoricalArrays include("callback.jl") using .CallBacks -function compute_grads(loss, model, batch) - dup_model = Duplicated(model) - const_args = map(Const, batch) - grads = Flux.gradient((m, args...) -> loss(m, args...), dup_model, const_args...) - return grads[1] -end - function init( config::LearnerTypes, df::AbstractDataFrame; @@ -144,6 +138,7 @@ function fit( while m.info[:nrounds] < config.nrounds fit_iter!(m, cache) + m.info[:nrounds] += 1 iter = m.info[:nrounds] if !isnothing(logger) cb(logger, iter, m) @@ -160,16 +155,32 @@ end function fit_iter!(m, cache) loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] - GC.gc(true) - if typeof(cache[:dtrain]) <: CUDA.CuIterator - CUDA.reclaim() - end for d in data - grads = compute_grads(loss, m, d) - Optimisers.update!(opts, m, grads) + const_args = map(Const, d) + grads = Enzyme.gradient(set_runtime_activity(Reverse), (m, args...) -> loss(m, args...), m, const_args...) + Optimisers.update!(opts, m, grads[1]) end - m.info[:nrounds] += 1 return nothing end +# function fit_iter!(m, cache) +# loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] +# GC.gc(true) +# if typeof(cache[:dtrain]) <: CUDA.CuIterator +# CUDA.reclaim() +# end +# for d in data +# grads = compute_grads(loss, m, d) +# Optimisers.update!(opts, m, grads) +# end +# m.info[:nrounds] += 1 +# return nothing +# end + +# function compute_grads(loss, model, batch) +# dup_model = Duplicated(model) +# const_args = map(Const, batch) +# grads = Flux.gradient((m, args...) -> loss(m, args...), dup_model, const_args...) +# return grads[1] +# end end \ No newline at end of file From 9e3d29d413d4b146678661cfd0da6f870f7fd5ab Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 31 Jan 2026 18:36:20 -0500 Subject: [PATCH 013/120] up --- src/Fit/fit.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 02069f1..e7eaac2 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -67,7 +67,6 @@ function init( optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) - cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) return m, cache end @@ -138,7 +137,6 @@ function fit( while m.info[:nrounds] < config.nrounds fit_iter!(m, cache) - m.info[:nrounds] += 1 iter = m.info[:nrounds] if !isnothing(logger) cb(logger, iter, m) @@ -160,6 +158,7 @@ function fit_iter!(m, cache) grads = Enzyme.gradient(set_runtime_activity(Reverse), (m, args...) -> loss(m, args...), m, const_args...) Optimisers.update!(opts, m, grads[1]) end + m.info[:nrounds] += 1 return nothing end # function fit_iter!(m, cache) From 020d4d02b80c026776461ed2aa35b26a7d5183ca Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 3 Feb 2026 21:24:15 -0500 Subject: [PATCH 014/120] fit return cpu model --- src/Fit/fit.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 8ceffa8..d1bdf10 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -152,8 +152,7 @@ function fit( end m.info[:logger] = logger - return m - # return m |> cpu + return m |> cpu end function fit_iter!(m, cache) From a342bd5eae70ef56fb9e1776c94faafdb047970a Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 5 Feb 2026 22:02:24 +0530 Subject: [PATCH 015/120] use reactant and flux approach Signed-off-by: AdityaPandeyCN --- Project.toml | 4 ++ benchmarks/titanic-logloss.jl | 3 ++ src/Fit/fit.jl | 56 +++++++++++++----------- src/data.jl | 2 +- src/infer.jl | 69 ++++++++++-------------------- src/models/NeuroTree/entmax.jl | 1 - src/models/NeuroTree/neurotrees.jl | 1 - 7 files changed, 62 insertions(+), 74 deletions(-) diff --git a/Project.toml b/Project.toml index 5065c25..f373770 100644 --- a/Project.toml +++ b/Project.toml @@ -4,16 +4,20 @@ authors = ["jeremie.db "] version = "0.3.0" [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 51f5f7d..ae31c07 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -8,6 +8,9 @@ using CategoricalArrays using OrderedCollections using NeuroTabModels +using Reactant +Reactant.set_default_backend("cpu") + Random.seed!(123) df = MLDatasets.Titanic().dataframe diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index d1bdf10..e0e3840 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -9,10 +9,12 @@ using ..Losses using ..Metrics import MLJModelInterface: fit -import CUDA, cuDNN +import Reactant +import Reactant: ConcreteRArray import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta -import Flux: trainmode!, gradient, cpu, gpu +import Optimisers: OptimiserChain, WeightDecay, Adam +import Flux +import ADTypes: AutoEnzyme using DataFrames using CategoricalArrays @@ -29,7 +31,6 @@ function init( offset_name=nothing, ) - device = config.device batchsize = config.batchsize nfeats = length(feature_names) loss = get_loss_fn(config.loss) @@ -47,7 +48,7 @@ function init( outsize = 2 end - dtrain = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) + dtrain = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) info = Dict( :nrounds => 0, @@ -57,11 +58,8 @@ function init( chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - if device == :gpu - m = m |> gpu - end - optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) @@ -82,8 +80,6 @@ end print_every_n=9999, early_stopping_rounds=9999, verbosity=1, - device=:cpu, - gpuID=0, ) Training function of NeuroTabModels' internal API. @@ -115,11 +111,6 @@ function fit( verbosity=1 ) - device = Symbol(config.device) - if device == :gpu - CUDA.device!(config.gpuID) - end - feature_names = Symbol.(feature_names) target_name = Symbol(target_name) weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) @@ -152,18 +143,35 @@ function fit( end m.info[:logger] = logger - return m |> cpu + return m end function fit_iter!(m, cache) - loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] - GC.gc(true) - if typeof(cache[:dtrain]) <: CUDA.CuIterator - CUDA.reclaim() - end + loss_fn, opts, data = cache[:loss], cache[:opts], cache[:dtrain] + for d in data - grads = gradient(model -> loss(model, d...), m)[1] - Optimisers.update!(opts, m, grads) + x, y = d[1], d[2] + w = length(d) >= 3 ? d[3] : nothing + o = length(d) >= 4 ? d[4] : nothing + + # Convert to Reactant arrays - backend (CPU/GPU) set via Reactant.set_default_backend() + x_ra = ConcreteRArray(x) + y_ra = ConcreteRArray(y) + m_ra = Reactant.to_rarray(m) + + # Dispatch based on available args - never pass nothing to traced code + if !isnothing(w) && !isnothing(o) + w_ra = ConcreteRArray(w) + o_ra = ConcreteRArray(o) + _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra, w_ra, o_ra) + elseif !isnothing(w) + w_ra = ConcreteRArray(w) + _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra, w_ra) + else + _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra) + end + + Optimisers.update!(opts, m, grads[1]) end m.info[:nrounds] += 1 return nothing diff --git a/src/data.jl b/src/data.jl index d41c2e8..4051a43 100644 --- a/src/data.jl +++ b/src/data.jl @@ -7,7 +7,7 @@ import MLUtils: DataLoader using DataFrames using CategoricalArrays -using CUDA: CuIterator + """ ContainerTrain diff --git a/src/infer.jl b/src/infer.jl index e09ec39..682ff25 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -4,107 +4,82 @@ using ..Data using ..Losses using ..Models -using Flux: sigmoid, softmax!, cpu, gpu, onecold +using Flux: sigmoid, softmax!, cpu using DataFrames: AbstractDataFrame import MLUtils: DataLoader -import CUDA: CuIterator, device! export infer """ - DL - -Union{NeuroTabModels.CuIterator, NeuroTabModels.DataLoader} -""" -const DL = Union{CuIterator,DataLoader} - -""" -infer(m::NeuroTabModel, data) - + infer(m::NeuroTabModel, data) Return the inference of a `NeuroTabModel` over `data`, where `data` is `AbstractDataFrame`. """ -function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu, gpuID=0) - if device == :gpu - device!(gpuID) - end - m = device == :gpu ? m |> gpu : m |> cpu - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048, device) +function infer(m::NeuroTabModel, data::AbstractDataFrame) + dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) p = infer(m, dinfer) return p end - -""" - (m::NeuroTabModel)(x::AbstractMatrix) - (m::NeuroTabModel)(data::AbstractDataFrame) - -Inference for NeuroTabModel -""" function (m::NeuroTabModel)(x::AbstractMatrix) p = m.chain(x) - if size(p, 1) == 1 - p = dropdims(p; dims=1) + if ndims(p) == 2 && size(p, 1) == 1 + p = vec(p) end return p end -function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu, gpuID=0) - if device == :gpu - device!(gpuID) - end - m = device == :gpu ? m |> gpu : m |> cpu - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048, device) + +function (m::NeuroTabModel)(data::AbstractDataFrame) + dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) p = infer(m, dinfer) return p end - -function infer(m::NeuroTabModel{L}, data::DL) where {L<:Union{MSE,MAE}} +function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{MSE,MAE}} preds = Vector{Float32}[] for x in data push!(preds, Vector(m(x))) end - p = vcat(preds...) + p = reduce(vcat, preds) return p end -function infer(m::NeuroTabModel{<:LogLoss}, data::DL) +function infer(m::NeuroTabModel{<:LogLoss}, data::DataLoader) preds = Vector{Float32}[] for x in data push!(preds, Vector(m(x))) end - p = vcat(preds...) - p .= sigmoid(p) + p = reduce(vcat, preds) + p = sigmoid.(p) return p end -function infer(m::NeuroTabModel{<:MLogLoss}, data::DL) +function infer(m::NeuroTabModel{<:MLogLoss}, data::DataLoader) preds = Matrix{Float32}[] for x in data push!(preds, Matrix(m(x)')) end - p = vcat(preds...) + p = reduce(vcat, preds) softmax!(p; dims=2) return p end -function infer(m::NeuroTabModel{<:GaussianMLE}, data::DL) +function infer(m::NeuroTabModel{<:GaussianMLE}, data::DataLoader) preds = Matrix{Float32}[] for x in data push!(preds, Matrix(m(x)')) end - p = vcat(preds...) - p[:, 2] .= exp.(p[:, 2]) # reproject log(σ) into σ + p = reduce(vcat, preds) return p end -function infer(m::NeuroTabModel{L}, data::DL) where {L<:Union{Tweedie}} +function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{Tweedie}} preds = Vector{Float32}[] for x in data push!(preds, Vector(m(x))) end - p = vcat(preds...) - p .= exp.(p) + p = reduce(vcat, preds) + p = exp.(p) return p end -end # module +end # module \ No newline at end of file diff --git a/src/models/NeuroTree/entmax.jl b/src/models/NeuroTree/entmax.jl index 4f4bc16..77dde49 100644 --- a/src/models/NeuroTree/entmax.jl +++ b/src/models/NeuroTree/entmax.jl @@ -1,6 +1,5 @@ # Reference: https://github.com/deep-spin/entmax/blob/master/entmax/activations.py one_to_vec(x) = reshape(vec(1:size(x, 2)), 1, :) -one_to_vec(x::AnyCuArray) = reshape(CUDA.CuArray(1:size(x, 2)), 1, :) function entmax_threshold_and_support(x) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 396e23b..0542371 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -3,7 +3,6 @@ module NeuroTrees export NeuroTreeConfig import .Threads: @threads -using CUDA import Flux import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu From 7c3d33e7ec4cead6163022c368ba2d512f26e034 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 6 Feb 2026 01:44:27 +0530 Subject: [PATCH 016/120] formatting Signed-off-by: AdityaPandeyCN --- Project.toml | 2 -- src/Fit/fit.jl | 8 ++++---- src/data.jl | 33 +++++++++------------------------ src/infer.jl | 19 ++++++++++--------- 4 files changed, 23 insertions(+), 39 deletions(-) diff --git a/Project.toml b/Project.toml index f373770..58cc4e6 100644 --- a/Project.toml +++ b/Project.toml @@ -11,11 +11,9 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" -MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" -OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index e0e3840..3103434 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -148,17 +148,17 @@ end function fit_iter!(m, cache) loss_fn, opts, data = cache[:loss], cache[:opts], cache[:dtrain] - + for d in data x, y = d[1], d[2] w = length(d) >= 3 ? d[3] : nothing o = length(d) >= 4 ? d[4] : nothing - + # Convert to Reactant arrays - backend (CPU/GPU) set via Reactant.set_default_backend() x_ra = ConcreteRArray(x) y_ra = ConcreteRArray(y) m_ra = Reactant.to_rarray(m) - + # Dispatch based on available args - never pass nothing to traced code if !isnothing(w) && !isnothing(o) w_ra = ConcreteRArray(w) @@ -170,7 +170,7 @@ function fit_iter!(m, cache) else _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra) end - + Optimisers.update!(opts, m, grads[1]) end m.info[:nrounds] += 1 diff --git a/src/data.jl b/src/data.jl index 4051a43..42d82c9 100644 --- a/src/data.jl +++ b/src/data.jl @@ -8,10 +8,8 @@ import MLUtils: DataLoader using DataFrames using CategoricalArrays - """ ContainerTrain - """ struct ContainerTrain{A<:AbstractMatrix,B<:AbstractVector,C,D} x::A @@ -48,7 +46,6 @@ function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B return (x, y, w, offset) end - function get_df_loader_train( df::AbstractDataFrame; feature_names, @@ -57,8 +54,7 @@ function get_df_loader_train( offset_name=nothing, batchsize, shuffle=true, - device=:cpu) - +) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') @@ -73,23 +69,17 @@ function get_df_loader_train( offset = if isnothing(offset_name) nothing else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, data.offset_name])') + isa(offset_name, String) ? Float32.(df[!, offset_name]) : Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') end container = ContainerTrain(x, y, w, offset) batchsize = min(batchsize, length(container)) dtrain = DataLoader(container; shuffle, batchsize, partial=true, parallel=false) - if device == :gpu - return CuIterator(dtrain) - else - return dtrain - end + return dtrain end - """ ContainerInfer - """ struct ContainerInfer{A<:AbstractMatrix,D} x::A @@ -102,12 +92,12 @@ function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:No x = data.x[:, idx] return x end -function getindex(data::ContainerTrain{A,D}, idx::AbstractVector) where {A,D<:AbstractVector} +function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:AbstractVector} x = data.x[:, idx] offset = data.offset[idx] return (x, offset) end -function getindex(data::ContainerTrain{A,D}, idx::AbstractVector) where {A,D<:AbstractMatrix} +function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:AbstractMatrix} x = data.x[:, idx] offset = data.offset[:, idx] return (x, offset) @@ -117,26 +107,21 @@ function get_df_loader_infer( df::AbstractDataFrame; feature_names, offset_name=nothing, - batchsize, - device=:cpu) - + batchsize=2048, +) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') offset = if isnothing(offset_name) nothing else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, data.offset_name])') + isa(offset_name, String) ? Float32.(df[!, offset_name]) : Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') end container = ContainerInfer(x, offset) batchsize = min(batchsize, length(container)) dinfer = DataLoader(container; shuffle=false, batchsize, partial=true, parallel=false) - if device == :gpu - return CuIterator(dinfer) - else - return dinfer - end + return dinfer end end #module \ No newline at end of file diff --git a/src/infer.jl b/src/infer.jl index 682ff25..1865f1f 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -22,8 +22,8 @@ end function (m::NeuroTabModel)(x::AbstractMatrix) p = m.chain(x) - if ndims(p) == 2 && size(p, 1) == 1 - p = vec(p) + if size(p, 1) == 1 + p = dropdims(p; dims=1) end return p end @@ -39,7 +39,7 @@ function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{MSE,MAE}} for x in data push!(preds, Vector(m(x))) end - p = reduce(vcat, preds) + p = vcat(preds...) return p end @@ -48,8 +48,8 @@ function infer(m::NeuroTabModel{<:LogLoss}, data::DataLoader) for x in data push!(preds, Vector(m(x))) end - p = reduce(vcat, preds) - p = sigmoid.(p) + p = vcat(preds...) + p .= sigmoid(p) return p end @@ -58,7 +58,7 @@ function infer(m::NeuroTabModel{<:MLogLoss}, data::DataLoader) for x in data push!(preds, Matrix(m(x)')) end - p = reduce(vcat, preds) + p = vcat(preds...) softmax!(p; dims=2) return p end @@ -68,7 +68,8 @@ function infer(m::NeuroTabModel{<:GaussianMLE}, data::DataLoader) for x in data push!(preds, Matrix(m(x)')) end - p = reduce(vcat, preds) + p = vcat(preds...) + p[:, 2] .= exp.(p[:, 2]) return p end @@ -77,8 +78,8 @@ function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{Tweedie}} for x in data push!(preds, Vector(m(x))) end - p = reduce(vcat, preds) - p = exp.(p) + p = vcat(preds...) + p .= exp.(p) return p end From c2ad582b179fbc48dbc420e4ebdafb7c2e600fc9 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 6 Feb 2026 10:07:27 +0530 Subject: [PATCH 017/120] revert changes Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 3103434..3c50b8c 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -59,7 +59,7 @@ function init( chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) From c0fc0ff4affe8e5aaf3846802d3156a504d1d88e Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 6 Feb 2026 10:15:26 +0530 Subject: [PATCH 018/120] add back imports Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 3c50b8c..55b4239 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -12,7 +12,7 @@ import MLJModelInterface: fit import Reactant import Reactant: ConcreteRArray import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam +import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta import Flux import ADTypes: AutoEnzyme From 01477acf9681b8c12553e8db0714bfb884bf6893 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 6 Feb 2026 10:56:16 +0530 Subject: [PATCH 019/120] update callback.jl Signed-off-by: AdityaPandeyCN --- src/Fit/callback.jl | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 276902a..23aa723 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -2,8 +2,7 @@ module CallBacks using DataFrames using Statistics: mean, median -using Flux: cpu, gpu -using CUDA: CuIterator +using Flux: cpu using ..Learners: LearnerTypes using ..Data: get_df_loader_train @@ -31,10 +30,9 @@ function CallBack( offset_name=nothing ) - device = config.device batchsize = config.batchsize feval = metric_dict[config.metric] - deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize, device) + deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize) return CallBack(feval, deval) end From 9c832e7399c2104d260a8860d243f338469c9054 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 6 Feb 2026 12:10:19 +0530 Subject: [PATCH 020/120] revert back NAdam and remove CUDA dependency Signed-off-by: AdityaPandeyCN --- Project.toml | 1 - src/Fit/fit.jl | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 58cc4e6..79bbe50 100644 --- a/Project.toml +++ b/Project.toml @@ -5,7 +5,6 @@ version = "0.3.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 55b4239..359247e 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -59,7 +59,7 @@ function init( chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) From 389d624bfe0280df20b95de76d42d4dbb9dda12b Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 9 Feb 2026 11:09:56 +0530 Subject: [PATCH 021/120] Compile gradient function once in init(), reuse across all epochs Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 56 +++++++++++++++++++++++++++++--------------------- 1 file changed, 33 insertions(+), 23 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 359247e..66feb92 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -12,7 +12,7 @@ import MLJModelInterface: fit import Reactant import Reactant: ConcreteRArray import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta +import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov import Flux import ADTypes: AutoEnzyme @@ -22,6 +22,18 @@ using CategoricalArrays include("callback.jl") using .CallBacks +function _compile_grad_fn(loss_fn, m, batch) + m_ra = Reactant.to_rarray(m) + args_ra = map(b -> ConcreteRArray(b), batch) + grad_fn = Reactant.@compile Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, args_ra...) + full_batchsize = size(batch[1], ndims(batch[1])) + return grad_fn, full_batchsize +end + +function _to_cpu_grads(grads) + return Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, grads) +end + function init( config::LearnerTypes, df::AbstractDataFrame; @@ -59,10 +71,17 @@ function init( chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - optim = OptimiserChain(Adam(config.lr), WeightDecay(config.wd)) + optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) - cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) + # Compile gradient function once using first batch as template + first_batch = first(dtrain) + grad_fn, full_batchsize = _compile_grad_fn(loss, m, first_batch) + + cache = ( + dtrain=dtrain, loss=loss, opts=opts, info=info, + grad_fn=grad_fn, full_batchsize=full_batchsize, + ) return m, cache end @@ -129,7 +148,6 @@ function fit( (verbosity > 0) && @info "Init training" end - # for iter = 1:config.nrounds while m.info[:nrounds] < config.nrounds fit_iter!(m, cache) iter = m.info[:nrounds] @@ -147,31 +165,23 @@ function fit( end function fit_iter!(m, cache) - loss_fn, opts, data = cache[:loss], cache[:opts], cache[:dtrain] - - for d in data - x, y = d[1], d[2] - w = length(d) >= 3 ? d[3] : nothing - o = length(d) >= 4 ? d[4] : nothing + loss_fn = cache[:loss] + opts = cache[:opts] + grad_fn = cache[:grad_fn] + full_batchsize = cache[:full_batchsize] - # Convert to Reactant arrays - backend (CPU/GPU) set via Reactant.set_default_backend() - x_ra = ConcreteRArray(x) - y_ra = ConcreteRArray(y) + for d in cache[:dtrain] m_ra = Reactant.to_rarray(m) + args_ra = map(b -> ConcreteRArray(b), d) + is_full = size(d[1], ndims(d[1])) == full_batchsize - # Dispatch based on available args - never pass nothing to traced code - if !isnothing(w) && !isnothing(o) - w_ra = ConcreteRArray(w) - o_ra = ConcreteRArray(o) - _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra, w_ra, o_ra) - elseif !isnothing(w) - w_ra = ConcreteRArray(w) - _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra, w_ra) + _, grads = if is_full + grad_fn(loss_fn, AutoEnzyme(), m_ra, args_ra...) else - _, grads = Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, x_ra, y_ra) + Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, args_ra...) end - Optimisers.update!(opts, m, grads[1]) + Optimisers.update!(opts, m, _to_cpu_grads(grads[1])) end m.info[:nrounds] += 1 return nothing From afcc48678496c6e1c964fc9330eec1b7fb046537 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 12 Feb 2026 23:44:29 +0530 Subject: [PATCH 022/120] remove CuDNN dependency Signed-off-by: AdityaPandeyCN --- Project.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/Project.toml b/Project.toml index 79bbe50..36f8076 100644 --- a/Project.toml +++ b/Project.toml @@ -18,10 +18,8 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] -CUDA = "5" CategoricalArrays = "1" ChainRulesCore = "1" DataFrames = "1.3" @@ -34,7 +32,6 @@ Random = "1" Statistics = "1" StatsBase = "0.34" Tables = "1.9" -cuDNN = "1" julia = "1.10" [extras] From 4258de8d11daef901171449dc57c3fb6ad938c31 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 12 Feb 2026 23:52:32 +0530 Subject: [PATCH 023/120] refactor fit.jl Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 98 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 66 insertions(+), 32 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 66feb92..0eaef29 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -10,11 +10,13 @@ using ..Metrics import MLJModelInterface: fit import Reactant -import Reactant: ConcreteRArray +import Reactant: ConcreteRArray, ConcreteRNumber import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov +import Optimisers: OptimiserChain, WeightDecay, NAdam import Flux +import Flux: onehotbatch import ADTypes: AutoEnzyme +import Functors: fmap using DataFrames using CategoricalArrays @@ -22,16 +24,37 @@ using CategoricalArrays include("callback.jl") using .CallBacks -function _compile_grad_fn(loss_fn, m, batch) - m_ra = Reactant.to_rarray(m) - args_ra = map(b -> ConcreteRArray(b), batch) - grad_fn = Reactant.@compile Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, args_ra...) - full_batchsize = size(batch[1], ndims(batch[1])) - return grad_fn, full_batchsize +function _to_rarray_with_scalars(x) + return fmap(x; exclude=v -> v isa AbstractArray{<:Number} || v isa Number) do v + if v isa AbstractArray{<:Number} + ConcreteRArray(v) + elseif v isa Bool + v + elseif v isa Number + ConcreteRNumber(Float32(v)) + else + v + end + end end -function _to_cpu_grads(grads) - return Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, grads) +function _train_step!(loss_fn, m, opts, args...) + _, grads = Flux.withgradient(loss_fn, AutoEnzyme(), m, args...) + new_opts, new_model = Optimisers.update!(opts, m, grads[1]) + return new_opts, new_model +end + +function _to_batches_ra(all_batches, L, outsize) + if L <: MLogLoss + return map(all_batches) do batch + x_ra = ConcreteRArray(batch[1]) + y_onehot = Float32.(onehotbatch(batch[2], UInt32(1):UInt32(outsize))) + y_ra = ConcreteRArray(y_onehot) + length(batch) > 2 ? (x_ra, y_ra, map(b -> ConcreteRArray(b), batch[3:end])...) : (x_ra, y_ra) + end + else + return [map(b -> ConcreteRArray(b), batch) for batch in all_batches] + end end function init( @@ -74,13 +97,25 @@ function init( optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) opts = Optimisers.setup(optim, m) - # Compile gradient function once using first batch as template - first_batch = first(dtrain) - grad_fn, full_batchsize = _compile_grad_fn(loss, m, first_batch) - - cache = ( - dtrain=dtrain, loss=loss, opts=opts, info=info, - grad_fn=grad_fn, full_batchsize=full_batchsize, + chain_ra = Reactant.to_rarray(m.chain) + m_ra = NeuroTabModel(m._loss_type, chain_ra, Dict{Symbol,Any}()) + opts_ra = _to_rarray_with_scalars(opts) + + all_batches = collect(dtrain) + batches_ra = _to_batches_ra(all_batches, L, outsize) + full_batchsize = size(all_batches[1][1], ndims(all_batches[1][1])) + is_full_batch = [size(b[1], ndims(b[1])) == full_batchsize for b in all_batches] + + compiled_step = Reactant.@compile _train_step!(loss, m_ra, opts_ra, batches_ra[1]...) + + cache = Dict( + :loss => loss, + :info => info, + :compiled_step => compiled_step, + :m_ra => m_ra, + :opts_ra => opts_ra, + :batches_ra => batches_ra, + :is_full_batch => is_full_batch ) return m, cache end @@ -152,6 +187,8 @@ function fit( fit_iter!(m, cache) iter = m.info[:nrounds] if !isnothing(logger) + chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) + Flux.loadmodel!(m.chain, chain_cpu) cb(logger, iter, m) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] @@ -160,29 +197,26 @@ function fit( end end + chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) + Flux.loadmodel!(m.chain, chain_cpu) m.info[:logger] = logger return m end function fit_iter!(m, cache) - loss_fn = cache[:loss] - opts = cache[:opts] - grad_fn = cache[:grad_fn] - full_batchsize = cache[:full_batchsize] - - for d in cache[:dtrain] - m_ra = Reactant.to_rarray(m) - args_ra = map(b -> ConcreteRArray(b), d) - is_full = size(d[1], ndims(d[1])) == full_batchsize - - _, grads = if is_full - grad_fn(loss_fn, AutoEnzyme(), m_ra, args_ra...) + m_ra = cache[:m_ra] + opts_ra = cache[:opts_ra] + + for (args_ra, is_full) in zip(cache[:batches_ra], cache[:is_full_batch]) + if is_full + opts_ra, m_ra = cache[:compiled_step](cache[:loss], m_ra, opts_ra, args_ra...) else - Reactant.@jit Flux.withgradient(loss_fn, AutoEnzyme(), m_ra, args_ra...) + opts_ra, m_ra = Reactant.@jit _train_step!(cache[:loss], m_ra, opts_ra, args_ra...) end - - Optimisers.update!(opts, m, _to_cpu_grads(grads[1])) end + + cache[:opts_ra] = opts_ra + cache[:m_ra] = m_ra m.info[:nrounds] += 1 return nothing end From 30897d0ba6bd5ac40256769a2661a7aacd22be04 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 12 Feb 2026 23:56:05 +0530 Subject: [PATCH 024/120] Add _sync_to_cpu! to MLJ interface to synchronize trained weights from Reactant cache Signed-off-by: AdityaPandeyCN --- src/MLJ.jl | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/src/MLJ.jl b/src/MLJ.jl index 4138738..8fbd5c1 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -4,11 +4,18 @@ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes import ..Fit: init, fit_iter! +import Reactant: ConcreteRArray +import Flux import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema export fit, update, predict +function _sync_to_cpu!(fitresult, cache) + chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) + Flux.loadmodel!(fitresult.chain, chain_cpu) +end + function fit( model::LearnerTypes, verbosity::Int, @@ -16,7 +23,8 @@ function fit( y, w=nothing) - Tables.istable(A) ? dtrain = DataFrame(A) : error("`A` must be a Table") + Tables.istable(A) || error("`A` must be a Table") + dtrain = DataFrame(A) feature_names = string.(collect(Tables.schema(dtrain).names)) @assert "_target" ∉ feature_names dtrain._target = y @@ -36,6 +44,7 @@ function fit( while fitresult.info[:nrounds] < model.nrounds fit_iter!(fitresult, cache) end + _sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) return fitresult, cache, report @@ -45,7 +54,6 @@ function okay_to_continue(model, fitresult, cache) return model.nrounds - fitresult.info[:nrounds] >= 0 end -# For EarlyStopping.jl support MMI.iteration_parameter(::Type{<:LearnerTypes}) = :nrounds function update( @@ -61,28 +69,29 @@ function update( while fitresult.info[:nrounds] < model.nrounds fit_iter!(fitresult, cache) end + _sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) + return fitresult, cache, report else fitresult, cache, report = fit(model, verbosity, A, y, w) + return fitresult, cache, report end - return fitresult, cache, report end -function predict(::NeuroTabRegressor, fitresult, A; device=:cpu, gpuID=0) +function predict(::NeuroTabRegressor, fitresult, A) + Tables.istable(A) || error("`A` must be a Table") df = DataFrame(A) - Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df; device, gpuID) + pred = fitresult(df) return pred end -function predict(::NeuroTabClassifier, fitresult, A; device=:cpu, gpuID=0) +function predict(::NeuroTabClassifier, fitresult, A) + Tables.istable(A) || error("`A` must be a Table") df = DataFrame(A) - Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df; device, gpuID) + pred = fitresult(df) return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end -# Metadata MMI.metadata_pkg.( (NeuroTabRegressor, NeuroTabClassifier), name="NeuroTabModels", @@ -109,4 +118,4 @@ MMI.metadata_model( path="NeuroTabModels.NeuroTabClassifier", ) -end +end \ No newline at end of file From 795027327fcd2bd42ccf35e1d0a209386880b8a0 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 12 Feb 2026 23:57:34 +0530 Subject: [PATCH 025/120] Update mlogloss to accept pre-encoded one-hot matrices, removing internal onehotbatch calls. Signed-off-by: AdityaPandeyCN --- src/losses.jl | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/losses.jl b/src/losses.jl index bec273e..7e645d2 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -4,7 +4,7 @@ export get_loss_fn, get_loss_type export LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE, Tweedie import Statistics: mean, std -import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid, onehotbatch +import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid abstract type LossType end abstract type MSE <: LossType end @@ -69,20 +69,19 @@ function tweedie(m, x, y, w, offset) ) / sum(w) end +# y is expected as a dense one-hot Float32 matrix (outsize, batchsize), +# pre-encoded in _to_batches_ra. No onehotbatch call needed here. function mlogloss(m, x, y) p = logsoftmax(m(x); dims=1) - k = size(p, 1) - mean(-sum(onehotbatch(y, 1:k) .* p; dims=1)) + mean(-sum(y .* p; dims=1)) end function mlogloss(m, x, y, w) p = logsoftmax(m(x); dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) + sum(-sum(y .* p; dims=1) .* w) / sum(w) end function mlogloss(m, x, y, w, offset) p = logsoftmax(m(x) .+ offset; dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) + sum(-sum(y .* p; dims=1) .* w) / sum(w) end gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}) where {T} = From 5587aede3dc756a0aba0aac3ea72b29e91097c46 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 12 Feb 2026 23:59:01 +0530 Subject: [PATCH 026/120] make track_stats false Signed-off-by: AdityaPandeyCN --- src/models/NeuroTree/neurotrees.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 0542371..9b230bd 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -79,7 +79,7 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) if config.MLE_tree_split && outsize == 2 outsize ÷= 2 chain = Chain( - BatchNorm(nfeats), + BatchNorm(nfeats, track_stats=false), Parallel( vcat, StackTree(nfeats => outsize; @@ -106,7 +106,7 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) ) else chain = Chain( - BatchNorm(nfeats), + BatchNorm(nfeats, track_stats=false), StackTree(nfeats => outsize; tree_type=config.tree_type, depth=config.depth, From 08974b42a1e514f73c980173561be9dbc326adb7 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 13 Feb 2026 00:00:05 +0530 Subject: [PATCH 027/120] update benchmark file Signed-off-by: AdityaPandeyCN --- benchmarks/benchmark_mse.jl | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index a4ba492..00bb85f 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -39,8 +39,7 @@ learner = NeuroTabRegressor( nrounds=10, early_stopping_rounds=2, lr=1e-2, - batchsize=2048, - device=:gpu + batchsize=2048 ) # desktop gpu: 13.476383 seconds (26.42 M allocations: 5.990 GiB, 9.44% gc time) @@ -54,4 +53,4 @@ learner = NeuroTabRegressor( ); # desktop: 0.771839 seconds (369.20 k allocations: 1.522 GiB, 5.94% gc time) -@time p_train = m(dtrain; device=:gpu); +@time p_train = m(dtrain); \ No newline at end of file From a0bfa18bbf0034210a095c4c964f99b35c440e92 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 13 Feb 2026 08:57:11 +0530 Subject: [PATCH 028/120] simplify cpu sync logic Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 14 ++++++++------ src/MLJ.jl | 9 +-------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 0eaef29..983f60b 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -1,6 +1,6 @@ module Fit -export fit +export fit, _sync_to_cpu! using ..Data using ..Learners @@ -57,6 +57,11 @@ function _to_batches_ra(all_batches, L, outsize) end end +function _sync_to_cpu!(m, cache) + chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) + Flux.loadmodel!(m.chain, chain_cpu) +end + function init( config::LearnerTypes, df::AbstractDataFrame; @@ -172,7 +177,6 @@ function fit( m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) - # initialize callback and logger if tracking eval data logger = nothing if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) @@ -187,8 +191,7 @@ function fit( fit_iter!(m, cache) iter = m.info[:nrounds] if !isnothing(logger) - chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) - Flux.loadmodel!(m.chain, chain_cpu) + _sync_to_cpu!(m, cache) cb(logger, iter, m) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] @@ -197,8 +200,7 @@ function fit( end end - chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) - Flux.loadmodel!(m.chain, chain_cpu) + _sync_to_cpu!(m, cache) m.info[:logger] = logger return m end diff --git a/src/MLJ.jl b/src/MLJ.jl index 8fbd5c1..51b5e99 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,19 +3,12 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit: init, fit_iter! -import Reactant: ConcreteRArray -import Flux +import ..Fit: init, fit_iter!, _sync_to_cpu! import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema export fit, update, predict -function _sync_to_cpu!(fitresult, cache) - chain_cpu = Flux.fmap(x -> x isa ConcreteRArray ? Array(x) : x, cache[:m_ra].chain) - Flux.loadmodel!(fitresult.chain, chain_cpu) -end - function fit( model::LearnerTypes, verbosity::Int, From c8b14d553f5d7daa4b3b55519a36ff9d7c9a4d9d Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 13 Feb 2026 09:19:53 +0530 Subject: [PATCH 029/120] clean MLJ.jl Signed-off-by: AdityaPandeyCN --- src/MLJ.jl | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/src/MLJ.jl b/src/MLJ.jl index 51b5e99..474b648 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -16,8 +16,7 @@ function fit( y, w=nothing) - Tables.istable(A) || error("`A` must be a Table") - dtrain = DataFrame(A) + Tables.istable(A) ? dtrain = DataFrame(A) : error("`A` must be a Table") feature_names = string.(collect(Tables.schema(dtrain).names)) @assert "_target" ∉ feature_names dtrain._target = y @@ -47,6 +46,7 @@ function okay_to_continue(model, fitresult, cache) return model.nrounds - fitresult.info[:nrounds] >= 0 end +# For EarlyStopping.jl support MMI.iteration_parameter(::Type{<:LearnerTypes}) = :nrounds function update( @@ -64,27 +64,25 @@ function update( end _sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) - return fitresult, cache, report else fitresult, cache, report = fit(model, verbosity, A, y, w) - return fitresult, cache, report end + return fitresult, cache, report end function predict(::NeuroTabRegressor, fitresult, A) - Tables.istable(A) || error("`A` must be a Table") - df = DataFrame(A) + Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") pred = fitresult(df) return pred end function predict(::NeuroTabClassifier, fitresult, A) - Tables.istable(A) || error("`A` must be a Table") - df = DataFrame(A) + Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") pred = fitresult(df) return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end +# Metadata MMI.metadata_pkg.( (NeuroTabRegressor, NeuroTabClassifier), name="NeuroTabModels", From 1066ecd66c597fbbb92ab56fa4f968fbda1ef73f Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 13 Feb 2026 12:22:12 +0530 Subject: [PATCH 030/120] fix imports Signed-off-by: AdityaPandeyCN --- src/MLJ.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/MLJ.jl b/src/MLJ.jl index 474b648..a45d558 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,7 +3,7 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit: init, fit_iter!, _sync_to_cpu! +import ..Fit: init, fit_iter! import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema @@ -36,7 +36,7 @@ function fit( while fitresult.info[:nrounds] < model.nrounds fit_iter!(fitresult, cache) end - _sync_to_cpu!(fitresult, cache) + Fit._sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) return fitresult, cache, report @@ -46,7 +46,6 @@ function okay_to_continue(model, fitresult, cache) return model.nrounds - fitresult.info[:nrounds] >= 0 end -# For EarlyStopping.jl support MMI.iteration_parameter(::Type{<:LearnerTypes}) = :nrounds function update( @@ -62,7 +61,7 @@ function update( while fitresult.info[:nrounds] < model.nrounds fit_iter!(fitresult, cache) end - _sync_to_cpu!(fitresult, cache) + Fit._sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) else fitresult, cache, report = fit(model, verbosity, A, y, w) @@ -82,7 +81,6 @@ function predict(::NeuroTabClassifier, fitresult, A) return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end -# Metadata MMI.metadata_pkg.( (NeuroTabRegressor, NeuroTabClassifier), name="NeuroTabModels", From 19f5c2567cd46fa04ddb25d9104ca8791d657664 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 13 Feb 2026 02:13:02 -0500 Subject: [PATCH 031/120] up --- benchmarks/benchmark_mse.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index a4ba492..accaf6b 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -6,7 +6,7 @@ using Random: seed! Threads.nthreads() seed!(123) -nobs = Int(1e6) +nobs = Int(1e5) num_feat = Int(100) @info "testing with: $nobs observations | $num_feat features." X = rand(Float32, nobs, num_feat) From 9a984a2fcd6b5f05d59e2eb2d3bb18cf3263ff98 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Fri, 13 Feb 2026 17:54:57 -0500 Subject: [PATCH 032/120] fix data loader --- src/data.jl | 8 ++++---- test/core.jl | 4 ++++ 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/data.jl b/src/data.jl index d41c2e8..d7e3573 100644 --- a/src/data.jl +++ b/src/data.jl @@ -73,7 +73,7 @@ function get_df_loader_train( offset = if isnothing(offset_name) nothing else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, data.offset_name])') + isa(offset_name, String) ? Float32.(df[!, offset_name]) : Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') end container = ContainerTrain(x, y, w, offset) @@ -102,12 +102,12 @@ function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:No x = data.x[:, idx] return x end -function getindex(data::ContainerTrain{A,D}, idx::AbstractVector) where {A,D<:AbstractVector} +function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:AbstractVector} x = data.x[:, idx] offset = data.offset[idx] return (x, offset) end -function getindex(data::ContainerTrain{A,D}, idx::AbstractVector) where {A,D<:AbstractMatrix} +function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:AbstractMatrix} x = data.x[:, idx] offset = data.offset[:, idx] return (x, offset) @@ -126,7 +126,7 @@ function get_df_loader_infer( offset = if isnothing(offset_name) nothing else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, data.offset_name])') + isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') end container = ContainerInfer(x, offset) diff --git a/test/core.jl b/test/core.jl index 07e7806..096f250 100644 --- a/test/core.jl +++ b/test/core.jl @@ -1,3 +1,7 @@ +@testset "Core - data iterators" begin + +end + @testset "Core - internals test" begin learner = NeuroTabRegressor(; From 94c667b6b936713e044e80458a5ad82d8f8f1cad Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 02:31:52 -0500 Subject: [PATCH 033/120] up --- Project.toml | 5 + benchmarks/titanic-logloss.jl | 34 ++----- experiments/dataloader.jl | 27 ++++-- experiments/lux.jl | 65 +++++++++++++ src/Fit/fit.jl | 67 ++++++++------ src/data.jl | 19 ++-- src/models/NeuroTree/model.jl | 143 ++++++++++------------------- src/models/NeuroTree/neurotrees.jl | 40 +++++--- src/models/models.jl | 13 ++- 9 files changed, 224 insertions(+), 189 deletions(-) create mode 100644 experiments/lux.jl diff --git a/Project.toml b/Project.toml index 5065c25..5e15b77 100644 --- a/Project.toml +++ b/Project.toml @@ -8,12 +8,17 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 51f5f7d..847a292 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -42,47 +42,27 @@ arch = NeuroTabModels.NeuroTreeConfig(; hidden_size=1, actA=:identity, ) -# arch = NeuroTabModels.MLPConfig(; -# act=:relu, -# stack_size=1, -# hidden_size=64, -# ) learner = NeuroTabRegressor( arch; loss=:logloss, - nrounds=200, + nrounds=100, early_stopping_rounds=2, lr=3e-2, device=:cpu ) -# learner = NeuroTabRegressor(; -# arch_name="NeuroTreeConfig", -# arch_config=Dict( -# :actA => :identity, -# :init_scale => 1.0, -# :depth => 4, -# :ntrees => 32, -# :stack_size => 1, -# :hidden_size => 1), -# loss=:logloss, -# nrounds=400, -# early_stopping_rounds=2, -# lr=1e-2, -# ) - @time m = NeuroTabModels.fit( learner, dtrain; - deval, + # deval, # removed - inference not yet adapted target_name, feature_names, print_every_n=10, ); -p_train = m(dtrain) -p_eval = m(deval) - -@info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) -@info mean((p_eval .> 0.5) .== (deval[!, target_name] .> 0.5)) +# commented out - inference not yet adapted +# p_train = m(dtrain) +# p_eval = m(deval) +# @info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) +# @info mean((p_eval .> 0.5) .== (deval[!, target_name] .> 0.5)) diff --git a/experiments/dataloader.jl b/experiments/dataloader.jl index b2352f1..b8d6f22 100644 --- a/experiments/dataloader.jl +++ b/experiments/dataloader.jl @@ -1,6 +1,7 @@ using NeuroTabModels using DataFrames using CategoricalArrays +using Lux ################################# # vanilla DataFrame @@ -9,7 +10,8 @@ nobs = 100 nfeats = 10 x = rand(nobs, nfeats); df = DataFrame(x, :auto); -df.y = rand(nobs); +y = rand(nobs); +df.y = y; target_name = "y" feature_names = Symbol.(setdiff(names(df), [target_name])) @@ -19,24 +21,35 @@ batchsize = 32 # CPU ################################### device = :cpu -dtrain = NeuroTabModels.get_df_loader_train(df; feature_names, target_name, batchsize, device) - +dtrain = NeuroTabModels.Data.get_df_loader_train(df; feature_names, target_name, batchsize, device) for d in dtrain @info length(d) - @info size(d[1]) + @info size(d[1]), size(d[2]) end -deval = NeuroTabModels.get_df_loader_infer(df; feature_names, batchsize=32) +deval = NeuroTabModels.Data.get_df_loader_infer(df; feature_names, batchsize=32) for d in deval @info size(d) end +################################### +# LuxDevice +################################### +# dev = reactant_device() +dev = cpu_device() +# dev = gpu_device() +dtrain = NeuroTabModels.Data.get_df_loader_train(df; feature_names, target_name, batchsize) |> dev +for d in dtrain + @info length(d) + @info size(d[1]) + @info typeof(d[1]) +end + ################################### # GPU ################################### device = :gpu -dtrain = NeuroTabModels.get_df_loader_train(df; feature_names, target_name, batchsize, device) - +dtrain = NeuroTabModels.Data.get_df_loader_train(df; feature_names, target_name, batchsize, device) for d in dtrain @info length(d) @info size(d[1]) diff --git a/experiments/lux.jl b/experiments/lux.jl new file mode 100644 index 0000000..c01b2c9 --- /dev/null +++ b/experiments/lux.jl @@ -0,0 +1,65 @@ +using NeuroTabModels +using Lux, LuxCore +using Random + +using NeuroTabModels.Models.NeuroTrees: get_logits_mask, get_softplus_mask +using NNlib: softplus +using Reactant +using Enzyme +using Optimisers + +rng = Random.Xoshiro(123) +nobs = 1000 +nfeats = 10 + +# m = NeuroTabModels.Models.NeuroTrees.NeuroTree(nfeats => 1; depth=4, trees=64) +m = Chain( + NeuroTabModels.Models.NeuroTrees.NeuroTree(nfeats => 1; depth=4, trees=64) +) +x = randn(rng, Float32, nfeats, nobs) +y = randn(rng, Float32, 1, nobs) +ps, st = LuxCore.setup(rng, m) +p, st = m(x, ps, st) + +# Get the device determined by Lux +# Reactant.set_default_backend("gpu") +Reactant.set_default_backend("gpu") +dev = reactant_device() +# dev = gpu_device() +# dev = cpu_device() + +# Parameter and State Variables +ps, st = Lux.setup(rng, m) |> dev + +# Dummy Input +x = rand(rng, Float32, nfeats, nobs) |> dev + +# Run the model +## We need to use @jit to compile and run the model with Reactant +@time p, st = @jit Lux.apply(m, x, ps, st) +# @time Lux.apply(m, x, ps, st) + +## For best performance, first compile the model with Reactant and then run it +@time apply_compiled = @compile Lux.apply(m, x, ps, st) +@time apply_compiled(m, x, ps, st) + +# Run the model +# Gradients +ts = Training.TrainState(m, ps, st, Adam(0.001f0)) +gs, loss, stats, ts = Lux.Training.compute_gradients( + AutoEnzyme(), + MSELoss(), + (x, y), + ts +) + +## Optimization +ts = Training.apply_gradients!(ts, gs) # or Training.apply_gradients (no `!` at the end) + +# Both these steps can be combined into a single call (preferred approach) +@time gs, loss, stats, ts = Training.single_train_step!( + AutoEnzyme(), + MSELoss(), + (x, y), + ts +); diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index d1bdf10..e98dccf 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -8,11 +8,14 @@ using ..Models using ..Losses using ..Metrics +import Random: Xoshiro import MLJModelInterface: fit import CUDA, cuDNN import Optimisers import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta -import Flux: trainmode!, gradient, cpu, gpu + +using Lux +using Enzyme, Reactant using DataFrames using CategoricalArrays @@ -47,7 +50,11 @@ function init( outsize = 2 end - dtrain = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) + Reactant.set_default_backend("gpu") + dev = reactant_device() + # dev = gpu_device() + # dev = cpu_device() + data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev info = Dict( :nrounds => 0, @@ -57,18 +64,17 @@ function init( chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - if device == :gpu - m = m |> gpu - end - optim = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) - opts = Optimisers.setup(optim, m) + # Parameter and State Variables + rng = Xoshiro(config.seed) + ps, st = Lux.setup(rng, m.chain) |> dev + opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + # ts = Training.TrainState(m.chain, ps, st, opt) - cache = (dtrain=dtrain, loss=loss, opts=opts, info=info) + cache = (data=data, ps=ps, st=st, opt=opt, info=info) return m, cache end - """ function fit( config::NeuroTypes, @@ -115,11 +121,6 @@ function fit( verbosity=1 ) - device = Symbol(config.device) - if device == :gpu - CUDA.device!(config.gpuID) - end - feature_names = Symbol.(feature_names) target_name = Symbol(target_name) weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) @@ -138,9 +139,22 @@ function fit( (verbosity > 0) && @info "Init training" end - # for iter = 1:config.nrounds + ts = Training.TrainState(m.chain, cache[:ps], cache[:st], cache[:opt]) while m.info[:nrounds] < config.nrounds - fit_iter!(m, cache) + # fit_iter!(m, cache[:ts], cache[:data]) + # ts, data = cache[:ts], cache[:data] + for d in cache[:data] + gs, loss, stats, ts = Training.single_train_step!( + AutoEnzyme(), + # MSELoss(), + BinaryCrossEntropyLoss(; logits=Val(true)), + (d[1], d[2]), + ts + ) + @info "loss: $loss" + end + m.info[:nrounds] += 1 + iter = m.info[:nrounds] if !isnothing(logger) cb(logger, iter, m) @@ -148,22 +162,23 @@ function fit( @info "iter $iter" metric = logger[:metrics][:metric][end] end (logger[:iter_since_best] >= logger[:early_stopping_rounds]) && break + else + (verbosity > 0 && iter % print_every_n == 0) && @info "iter $iter" end end - m.info[:logger] = logger - return m |> cpu + return m end -function fit_iter!(m, cache) - loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] - GC.gc(true) - if typeof(cache[:dtrain]) <: CUDA.CuIterator - CUDA.reclaim() - end +function fit_iter!(m, ts, data) + # ts, data = cache[:ts], cache[:data] for d in data - grads = gradient(model -> loss(model, d...), m)[1] - Optimisers.update!(opts, m, grads) + gs, loss, stats, ts = Training.single_train_step!( + AutoEnzyme(), + MSELoss(), + (d[1], d[2]), + ts + ) end m.info[:nrounds] += 1 return nothing diff --git a/src/data.jl b/src/data.jl index d7e3573..7fdd5da 100644 --- a/src/data.jl +++ b/src/data.jl @@ -13,7 +13,7 @@ using CUDA: CuIterator ContainerTrain """ -struct ContainerTrain{A<:AbstractMatrix,B<:AbstractVector,C,D} +struct ContainerTrain{A<:AbstractMatrix,B,C,D} x::A y::B w::C @@ -24,25 +24,25 @@ length(data::ContainerTrain) = size(data.x, 2) function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B,C<:Nothing,D<:Nothing} x = data.x[:, idx] - y = data.y[idx] + y = data.y[1:1, idx] return (x, y) end function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B,C<:AbstractVector,D<:Nothing} x = data.x[:, idx] - y = data.y[idx] + y = data.y[1:1, idx] w = data.w[idx] return (x, y, w) end function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B,C<:AbstractVector,D<:AbstractVector} x = data.x[:, idx] - y = data.y[idx] + y = data.y[1:1, idx] w = data.w[idx] offset = data.offset[idx] return (x, y, w, offset) end function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B,C<:AbstractVector,D<:AbstractMatrix} x = data.x[:, idx] - y = data.y[idx] + y = data.y[1:1, idx] w = data.w[idx] offset = data.offset[:, idx] return (x, y, w, offset) @@ -67,6 +67,7 @@ function get_df_loader_train( else y = Float32.(df[!, target_name]) end + y = reshape(y, 1, :) w = isnothing(weight_name) ? nothing : Float32.(df[!, weight_name]) @@ -78,12 +79,8 @@ function get_df_loader_train( container = ContainerTrain(x, y, w, offset) batchsize = min(batchsize, length(container)) - dtrain = DataLoader(container; shuffle, batchsize, partial=true, parallel=false) - if device == :gpu - return CuIterator(dtrain) - else - return dtrain - end + dtrain = DataLoader(container; shuffle, batchsize, partial=false, parallel=false) + return dtrain end diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 9d03e00..23b3fa2 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -1,77 +1,55 @@ -struct NeuroTree{M,V,F} - w::M - b::V - s::V - p::M - ml::M - ms::M +struct NeuroTree{F} <: AbstractLuxLayer + tree_type::Symbol actA::F scaler::Bool - ntrees::Int + feats::Int + outs::Int + depth::Int + trees::Int + nodes::Int + leaves::Int + init_scale::Float32 end -@layer NeuroTree trainable = (w, b, s, p) -# h = m.d_in(x) # [F,B] => [HNT,B] -# h = reshape(h, size(m.d_proj.weight, 2), :) # [HNT,B] => [H,NTB] -# nw = m.d_proj(h) # [H,NTB] => [1,NTB] -# nw = reshape(nw, size(m.lmask, 2), :) # [1,NTB] => [N,TB] -function (m::NeuroTree)(x) - nw = m.w * x .+ m.b # [F,B] => [NT,B] - if m.scaler - nw = softplus(m.s) .* (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] - else - nw = (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] - end - nw = reshape(nw, size(m.ml, 2), :) # [NT,B] => [N,TB] - lw = exp.(m.ml * nw .- m.ms * softplus.(nw)) # [N,TB] => [L,TB] - lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] - p = m.p * lw ./ m.ntrees # [P,LT] * [LT,B] => [P,B] - return p +function NeuroTree(; feats, outs, tree_type=:binary, actA=identity, scaler=true, depth, trees, init_scale=0.1) + nodes = 2^depth - 1 + leaves = 2^depth + return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, Float32(init_scale)) +end +function NeuroTree((feats, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, actA=identity, scaler=true, depth, trees, init_scale=0.1) + nodes = 2^depth - 1 + leaves = 2^depth + return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, Float32(init_scale)) end -""" - NeuroTree(; ins, outs, depth=4, ntrees=64, actA=identity, init_scale=1e-1) - NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; depth=4, ntrees=64, actA=identity, init_scale=1e-1) - -Initialization of a NeuroTree. -""" -function NeuroTree(; ins, outs, tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) - ml = get_logits_mask(Val(tree_type), depth) - ms = get_softplus_mask(Val(tree_type), depth) - nleaves = size(ml, 1) - nnodes = size(ml, 2) - - op = NeuroTree( - Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - zeros(Float32, nnodes * ntrees), # b - Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s - Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p - Float32.(ml), - Float32.(ms), - actA, - scaler, - ntrees, +# Define the Lux interface +function LuxCore.initialparameters(rng::AbstractRNG, l::NeuroTree) + return ( + w=Float32.((rand(l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # w + b=zeros(Float32, l.nodes * l.trees), # b + s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), # s + p=Float32.(randn(l.outs, l.leaves * l.trees) .* l.init_scale), # p ) - return op end -function NeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, actA=identity, scaler=true, init_scale=1e-1) - ml = get_logits_mask(Val(tree_type), depth) - ms = get_softplus_mask(Val(tree_type), depth) - nleaves = size(ml, 1) - nnodes = size(ml, 2) - op = NeuroTree( - Float32.((rand(nnodes * ntrees, ins) .- 0.5) ./ 4), # w - Float32.((rand(nnodes * ntrees) .- 0.5) ./ 4), # b - Float32.(fill(log(exp(1) - 1), nnodes * ntrees)), # s - Float32.(randn(outs, nleaves * ntrees) .* init_scale), # p - Float32.(ml), - Float32.(ms), - actA, - scaler, - ntrees, +function LuxCore.initialstates(rng::AbstractRNG, l::NeuroTree) + return ( + ml=get_logits_mask(Val(l.tree_type), l.depth), + ms=get_softplus_mask(Val(l.tree_type), l.depth) ) - return op +end + +function (l::NeuroTree)(x, ps, st) + if l.scaler + nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] + else + nw = (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] + end + nw = reshape(nw, size(st.ml, 2), :) # [NT,B] => [N,TB] + lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) # [N,TB] => [L,TB] + lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] + y = ps.p * lw ./ l.trees # [P,LT] * [LT,B] => [P,B] + return y, st end @@ -129,34 +107,6 @@ function get_softplus_mask(::Val{:oblivious}, depth::Integer) return mask end -function get_mask(::Val{:binary}, depth::Integer) - nodes = 2^depth - 1 - leaves = 2^depth - mask = zeros(Float32, nodes, leaves) - for d in 1:depth - blocks = 2^(d - 1) - k = 2^(depth - d) - stride = 2 * k - for b in 1:blocks - view(mask, 2^(d - 1) + b - 1, (b-1)*stride+1:(b-1)*stride+k) .= 1 - end - end - return mask -end -function get_mask(::Val{:oblivious}, depth::Integer) - leaves = 2^depth - mask = zeros(Bool, depth, leaves) - for d in 1:depth - blocks = 2^(d - 1) - k = 2^(depth - d) - stride = 2 * k - for b in 1:blocks - view(mask, d, (b-1)*stride+1:(b-1)*stride+k) .= true - end - end - return mask -end - """ StackTree @@ -165,7 +115,6 @@ A StackTree is made of a collection of NeuroTree. struct StackTree trees::Vector{NeuroTree} end -@layer StackTree function StackTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, stack_size=1, hidden_size=8, actA=identity, scaler=true, init_scale=1e-1) @assert stack_size == 1 || hidden_size >= outs @@ -173,17 +122,17 @@ function StackTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, de for i in 1:stack_size if i == 1 if i < stack_size - tree = NeuroTree(ins => hidden_size; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) + tree = NeuroTree(ins => hidden_size; tree_type, depth, trees, proj_size, actA, scaler, init_scale) push!(trees, tree) else - tree = NeuroTree(ins => outs; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) + tree = NeuroTree(ins => outs; tree_type, depth, trees, proj_size, actA, scaler, init_scale) push!(trees, tree) end elseif i < stack_size - tree = NeuroTree(hidden_size => hidden_size; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) + tree = NeuroTree(hidden_size => hidden_size; tree_type, depth, trees, proj_size, actA, scaler, init_scale) push!(trees, tree) else - tree = NeuroTree(hidden_size => outs; tree_type, depth, ntrees, proj_size, actA, scaler, init_scale) + tree = NeuroTree(hidden_size => outs; tree_type, depth, trees, proj_size, actA, scaler, init_scale) push!(trees, tree) end end diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 396e23b..f2d233a 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -2,13 +2,17 @@ module NeuroTrees export NeuroTreeConfig +using Random import .Threads: @threads -using CUDA -import Flux -import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch, glorot_uniform -import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel +# using CUDA +# import Flux +# import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu +# import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch, glorot_uniform +# import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel + +using Lux +using NNlib: softplus, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture @@ -106,20 +110,29 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) ) ) else + # chain = Chain( + # BatchNorm(nfeats), + # StackTree(nfeats => outsize; + # tree_type=config.tree_type, + # depth=config.depth, + # ntrees=config.ntrees, + # proj_size=config.proj_size, + # stack_size=config.stack_size, + # hidden_size=config.hidden_size, + # actA=act_dict[config.actA], + # scaler=config.scaler, + # init_scale=config.init_scale) + # ) chain = Chain( - BatchNorm(nfeats), - StackTree(nfeats => outsize; + # BatchNorm(nfeats), + NeuroTree(nfeats => outsize; tree_type=config.tree_type, depth=config.depth, - ntrees=config.ntrees, - proj_size=config.proj_size, - stack_size=config.stack_size, - hidden_size=config.hidden_size, + trees=config.ntrees, actA=act_dict[config.actA], scaler=config.scaler, - init_scale=config.init_scale) + init_scale=config.init_scale), ) - end end @@ -151,5 +164,4 @@ const act_dict = Dict( :hardtanh => _hardtanh_act, ) - end \ No newline at end of file diff --git a/src/models/models.jl b/src/models/models.jl index 41f3ddd..61033ac 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -4,8 +4,7 @@ export NeuroTabModel, Architecture export NeuroTreeConfig, MLPConfig, ResNetConfig using ..Losses -import Flux: @layer, Chain -import Functors: @functor +using Lux: Chain abstract type Architecture end @@ -17,13 +16,13 @@ struct NeuroTabModel{L<:LossType,C<:Chain} chain::C info::Dict{Symbol,Any} end -@functor NeuroTabModel (chain,) +# @functor NeuroTabModel (chain,) include("NeuroTree/neurotrees.jl") using .NeuroTrees -include("MLP/mlp.jl") -using .MLP -include("ResNet/resnet.jl") -using .ResNet +# include("MLP/mlp.jl") +# using .MLP +# include("ResNet/resnet.jl") +# using .ResNet end \ No newline at end of file From 1acb975e9426308e499d0cf6b213d4eb48ef3dfb Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 02:34:45 -0500 Subject: [PATCH 034/120] up --- src/Fit/fit.jl | 36 ++++++++++++++++-------------------- 1 file changed, 16 insertions(+), 20 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index e98dccf..f3cd907 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -50,7 +50,7 @@ function init( outsize = 2 end - Reactant.set_default_backend("gpu") + Reactant.set_default_backend("cpu") dev = reactant_device() # dev = gpu_device() # dev = cpu_device() @@ -69,7 +69,6 @@ function init( rng = Xoshiro(config.seed) ps, st = Lux.setup(rng, m.chain) |> dev opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) - # ts = Training.TrainState(m.chain, ps, st, opt) cache = (data=data, ps=ps, st=st, opt=opt, info=info) return m, cache @@ -141,17 +140,14 @@ function fit( ts = Training.TrainState(m.chain, cache[:ps], cache[:st], cache[:opt]) while m.info[:nrounds] < config.nrounds - # fit_iter!(m, cache[:ts], cache[:data]) - # ts, data = cache[:ts], cache[:data] for d in cache[:data] gs, loss, stats, ts = Training.single_train_step!( AutoEnzyme(), - # MSELoss(), - BinaryCrossEntropyLoss(; logits=Val(true)), + MSELoss(), + # BinaryCrossEntropyLoss(; logits=Val(true)), (d[1], d[2]), ts ) - @info "loss: $loss" end m.info[:nrounds] += 1 @@ -170,18 +166,18 @@ function fit( return m end -function fit_iter!(m, ts, data) - # ts, data = cache[:ts], cache[:data] - for d in data - gs, loss, stats, ts = Training.single_train_step!( - AutoEnzyme(), - MSELoss(), - (d[1], d[2]), - ts - ) - end - m.info[:nrounds] += 1 - return nothing -end +# function fit_iter!(m, ts, data) +# # ts, data = cache[:ts], cache[:data] +# for d in data +# gs, loss, stats, ts = Training.single_train_step!( +# AutoEnzyme(), +# MSELoss(), +# (d[1], d[2]), +# ts +# ) +# end +# m.info[:nrounds] += 1 +# return nothing +# end end \ No newline at end of file From 5c6f97bfad958f300431fd277a8b2f45c0b26aa3 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 02:49:06 -0500 Subject: [PATCH 035/120] up --- src/Fit/fit.jl | 3 ++- src/models/NeuroTree/neurotrees.jl | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index f3cd907..b79f0b2 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -50,8 +50,9 @@ function init( outsize = 2 end - Reactant.set_default_backend("cpu") + Reactant.set_default_backend("gpu") dev = reactant_device() + @info "dev" dev # dev = gpu_device() # dev = cpu_device() data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index f2d233a..8db3ce1 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -124,7 +124,7 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) # init_scale=config.init_scale) # ) chain = Chain( - # BatchNorm(nfeats), + BatchNorm(nfeats), NeuroTree(nfeats => outsize; tree_type=config.tree_type, depth=config.depth, From b8becdff6ab4aad18e72375984e3d98ceb65f0ba Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 03:11:15 -0500 Subject: [PATCH 036/120] up --- src/Fit/fit.jl | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index b79f0b2..30fd06e 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -52,9 +52,8 @@ function init( Reactant.set_default_backend("gpu") dev = reactant_device() - @info "dev" dev - # dev = gpu_device() # dev = cpu_device() + # dev = gpu_device() data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev info = Dict( @@ -167,18 +166,4 @@ function fit( return m end -# function fit_iter!(m, ts, data) -# # ts, data = cache[:ts], cache[:data] -# for d in data -# gs, loss, stats, ts = Training.single_train_step!( -# AutoEnzyme(), -# MSELoss(), -# (d[1], d[2]), -# ts -# ) -# end -# m.info[:nrounds] += 1 -# return nothing -# end - end \ No newline at end of file From 8c02e99edf3265fb63887f2397c8528b5857d663 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 03:14:09 -0500 Subject: [PATCH 037/120] up --- benchmarks/benchmark_mse.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index accaf6b..09c8d0a 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -54,4 +54,5 @@ learner = NeuroTabRegressor( ); # desktop: 0.771839 seconds (369.20 k allocations: 1.522 GiB, 5.94% gc time) -@time p_train = m(dtrain; device=:gpu); +# FIXME: need to adapt infer +# @time p_train = m(dtrain; device=:gpu); From 47bd880a94f9be730f1990b6e4161f89e4da9311 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 03:25:21 -0500 Subject: [PATCH 038/120] up --- benchmarks/benchmark_mse.jl | 2 +- src/Fit/fit.jl | 35 +++++++++++++++++++++-------------- src/models/NeuroTree/model.jl | 1 - 3 files changed, 22 insertions(+), 16 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index accaf6b..a4ba492 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -6,7 +6,7 @@ using Random: seed! Threads.nthreads() seed!(123) -nobs = Int(1e5) +nobs = Int(1e6) num_feat = Int(100) @info "testing with: $nobs observations | $num_feat features." X = rand(Float32, nobs, num_feat) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index d1bdf10..f6856a5 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -138,10 +138,21 @@ function fit( (verbosity > 0) && @info "Init training" end + GC.gc(true) + if typeof(cache[:dtrain]) <: CUDA.CuIterator + CUDA.reclaim() + end # for iter = 1:config.nrounds + loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] while m.info[:nrounds] < config.nrounds - fit_iter!(m, cache) + # fit_iter!(m, cache) + for d in data + grads = gradient(model -> loss(model, d...), m)[1] + Optimisers.update!(opts, m, grads) + end + m.info[:nrounds] += 1 iter = m.info[:nrounds] + if !isnothing(logger) cb(logger, iter, m) if verbosity > 0 && iter % print_every_n == 0 @@ -155,18 +166,14 @@ function fit( return m |> cpu end -function fit_iter!(m, cache) - loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] - GC.gc(true) - if typeof(cache[:dtrain]) <: CUDA.CuIterator - CUDA.reclaim() - end - for d in data - grads = gradient(model -> loss(model, d...), m)[1] - Optimisers.update!(opts, m, grads) - end - m.info[:nrounds] += 1 - return nothing -end +# function fit_iter!(m, cache) +# loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] +# for d in data +# grads = gradient(model -> loss(model, d...), m)[1] +# Optimisers.update!(opts, m, grads) +# end +# m.info[:nrounds] += 1 +# return nothing +# end end \ No newline at end of file diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 9d03e00..a945b5c 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -16,7 +16,6 @@ end # nw = m.d_proj(h) # [H,NTB] => [1,NTB] # nw = reshape(nw, size(m.lmask, 2), :) # [1,NTB] => [N,TB] function (m::NeuroTree)(x) - nw = m.w * x .+ m.b # [F,B] => [NT,B] if m.scaler nw = softplus(m.s) .* (m.actA(m.w) * x .+ m.b) # [F,B] => [NT,B] else From e8a28e71c1fe3bd0b8c7fffa6dc4c7f7e29fefe6 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 03:32:16 -0500 Subject: [PATCH 039/120] up --- src/Fit/fit.jl | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index f6856a5..9247abb 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -145,12 +145,7 @@ function fit( # for iter = 1:config.nrounds loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] while m.info[:nrounds] < config.nrounds - # fit_iter!(m, cache) - for d in data - grads = gradient(model -> loss(model, d...), m)[1] - Optimisers.update!(opts, m, grads) - end - m.info[:nrounds] += 1 + fit_iter!(m, cache) iter = m.info[:nrounds] if !isnothing(logger) @@ -166,14 +161,14 @@ function fit( return m |> cpu end -# function fit_iter!(m, cache) -# loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] -# for d in data -# grads = gradient(model -> loss(model, d...), m)[1] -# Optimisers.update!(opts, m, grads) -# end -# m.info[:nrounds] += 1 -# return nothing -# end +function fit_iter!(m, cache) + loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] + for d in data + grads = gradient(model -> loss(model, d...), m)[1] + Optimisers.update!(opts, m, grads) + end + m.info[:nrounds] += 1 + return nothing +end end \ No newline at end of file From 3ca7a83a0e0fded87c4e960da8932085ebe3aa79 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sun, 15 Feb 2026 03:43:41 -0500 Subject: [PATCH 040/120] up --- benchmarks/benchmark_mse.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 09c8d0a..8988d2a 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -6,7 +6,7 @@ using Random: seed! Threads.nthreads() seed!(123) -nobs = Int(1e5) +nobs = Int(1e6) num_feat = Int(100) @info "testing with: $nobs observations | $num_feat features." X = rand(Float32, nobs, num_feat) From c7fa43a68d9d95187e27e83b1f0bd977aed3bed6 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sun, 15 Feb 2026 17:55:59 +0530 Subject: [PATCH 041/120] fix inference, callbacks, and classification Signed-off-by: AdityaPandeyCN Signed-off-by: AdityaPandeyCN --- Project.toml | 14 ++--- src/Fit/callback.jl | 11 +--- src/Fit/fit.jl | 124 +++++++++++++++++++++++++----------- src/MLJ.jl | 149 ++++++++++++++++++++++---------------------- src/data.jl | 19 ++---- src/infer.jl | 76 ++++++++++------------ src/losses.jl | 15 ++--- src/metrics.jl | 16 ++--- 8 files changed, 220 insertions(+), 204 deletions(-) diff --git a/Project.toml b/Project.toml index 5e15b77..23f4c6b 100644 --- a/Project.toml +++ b/Project.toml @@ -4,44 +4,38 @@ authors = ["jeremie.db "] version = "0.3.0" [deps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" -Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] -CUDA = "5" CategoricalArrays = "1" -ChainRulesCore = "1" DataFrames = "1.3" -Flux = "0.14, 0.15, 0.16" -Functors = "0.5" +Lux = "1" MLJModelInterface = "1.2.1" MLUtils = "0.4" +NNlib = "0.9" Optimisers = "0.4" Random = "1" Statistics = "1" StatsBase = "0.34" Tables = "1.9" -cuDNN = "1" julia = "1.10" [extras] +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d" MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 276902a..c063989 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -2,8 +2,6 @@ module CallBacks using DataFrames using Statistics: mean, median -using Flux: cpu, gpu -using CUDA: CuIterator using ..Learners: LearnerTypes using ..Data: get_df_loader_train @@ -30,11 +28,9 @@ function CallBack( weight_name=nothing, offset_name=nothing ) - - device = config.device batchsize = config.batchsize feval = metric_dict[config.metric] - deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize, device) + deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize) return CallBack(feval, deval) end @@ -71,14 +67,11 @@ function update_logger!(logger; iter, metric) end function agg_logger(logger_raw::Vector{Dict}) - _l1 = first(logger_raw) best_iters = [d[:best_iter] for d in logger_raw] best_iter = ceil(Int, median(best_iters)) - best_metrics = [d[:best_metric] for d in logger_raw] best_metric = last(best_metrics) - metrics = (layer=Int[], iter=Int[], metric=Float64[]) for i in eachindex(logger_raw) _l = logger_raw[i] @@ -86,7 +79,6 @@ function agg_logger(logger_raw::Vector{Dict}) append!(metrics[:iter], _l[:metrics][:iter]) append!(metrics[:metric], _l[:metrics][:metric]) end - logger = Dict( :name => _l1[:name], :maximise => _l1[:maximise], @@ -97,7 +89,6 @@ function agg_logger(logger_raw::Vector{Dict}) :best_metrics => best_metrics, :best_metric => best_metric, ) - return logger end diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 30fd06e..23dda8e 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -1,6 +1,6 @@ module Fit -export fit +export fit, fit_iter! using ..Data using ..Learners @@ -10,12 +10,12 @@ using ..Metrics import Random: Xoshiro import MLJModelInterface: fit -import CUDA, cuDNN -import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta +import Optimisers: OptimiserChain, WeightDecay, NAdam +import OneHotArrays: onehotbatch using Lux -using Enzyme, Reactant +using Reactant +using Lux: cpu_device, reactant_device using DataFrames using CategoricalArrays @@ -23,6 +23,37 @@ using CategoricalArrays include("callback.jl") using .CallBacks +function _get_device(config) + device = Symbol(config.device) + if device == :gpu + try + Reactant.set_default_backend("gpu") + return reactant_device() + catch + @warn "GPU not available, falling back to CPU" + Reactant.set_default_backend("cpu") + return reactant_device() + end + else + Reactant.set_default_backend("cpu") + return reactant_device() + end +end + +function _get_lux_loss(L) + if L <: MSE + return MSELoss() + elseif L <: MAE + return MAELoss() + elseif L <: LogLoss + return BinaryCrossEntropyLoss(; logits=Val(true)) + elseif L <: MLogLoss + return CrossEntropyLoss(; logits=Val(true)) + else + return MSELoss() + end +end + function init( config::LearnerTypes, df::AbstractDataFrame; @@ -31,12 +62,12 @@ function init( weight_name=nothing, offset_name=nothing, ) + dev = _get_device(config) - device = config.device batchsize = config.batchsize nfeats = length(feature_names) - loss = get_loss_fn(config.loss) L = get_loss_type(config.loss) + lux_loss = _get_lux_loss(L) target_levels = nothing target_isordered = false @@ -50,11 +81,19 @@ function init( outsize = 2 end - Reactant.set_default_backend("gpu") - dev = reactant_device() - # dev = cpu_device() - # dev = gpu_device() - data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev + dtrain_loader = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) + all_batches = collect(dtrain_loader) + + # One-hot encode targets for multiclass classification + if L <: MLogLoss + all_batches = map(all_batches) do batch + x = batch[1] + y_oh = Float32.(onehotbatch(vec(batch[2]), UInt32(1):UInt32(outsize))) + length(batch) > 2 ? (x, y_oh, batch[3:end]...) : (x, y_oh) + end + end + + data = all_batches |> dev info = Dict( :nrounds => 0, @@ -65,30 +104,30 @@ function init( chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - # Parameter and State Variables rng = Xoshiro(config.seed) ps, st = Lux.setup(rng, m.chain) |> dev opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + ts = Training.TrainState(m.chain, ps, st, opt) - cache = (data=data, ps=ps, st=st, opt=opt, info=info) + cache = Dict( + :data => data, + :lux_loss => lux_loss, + :train_state => ts, + ) return m, cache end """ function fit( - config::NeuroTypes, + config::LearnerTypes, dtrain; feature_names, target_name, weight_name=nothing, offset_name=nothing, deval=nothing, - metric=nothing, print_every_n=9999, - early_stopping_rounds=9999, verbosity=1, - device=:cpu, - gpuID=0, ) Training function of NeuroTabModels' internal API. @@ -96,15 +135,15 @@ Training function of NeuroTabModels' internal API. # Arguments - `config::LearnerTypes` -- `dtrain`: Must be `<:AbstractDataFrame` +- `dtrain`: Must be `<:AbstractDataFrame` # Keyword arguments - `feature_names`: Required kwarg, a `Vector{Symbol}` or `Vector{String}` of the feature names. -- `target_name` Required kwarg, a `Symbol` or `String` indicating the name of the target variable. +- `target_name`: Required kwarg, a `Symbol` or `String` indicating the name of the target variable. - `weight_name=nothing` - `offset_name=nothing` -- `deval=nothing` Data for tracking evaluation metric and perform early stopping. +- `deval=nothing`: Data for tracking evaluation metric and perform early stopping. - `print_every_n=9999` - `verbosity=1` """ @@ -127,32 +166,23 @@ function fit( m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) - # initialize callback and logger if tracking eval data logger = nothing if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) + _sync_params_to_model!(m, cache) cb(logger, 0, m) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] else (verbosity > 0) && @info "Init training" end - ts = Training.TrainState(m.chain, cache[:ps], cache[:st], cache[:opt]) while m.info[:nrounds] < config.nrounds - for d in cache[:data] - gs, loss, stats, ts = Training.single_train_step!( - AutoEnzyme(), - MSELoss(), - # BinaryCrossEntropyLoss(; logits=Val(true)), - (d[1], d[2]), - ts - ) - end - m.info[:nrounds] += 1 - + fit_iter!(m, cache) iter = m.info[:nrounds] + if !isnothing(logger) + _sync_params_to_model!(m, cache) cb(logger, iter, m) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] @@ -162,8 +192,32 @@ function fit( (verbosity > 0 && iter % print_every_n == 0) && @info "iter $iter" end end + + _sync_params_to_model!(m, cache) m.info[:logger] = logger return m end +function _sync_params_to_model!(m, cache) + ts = cache[:train_state] + cdev = cpu_device() + m.info[:ps] = cdev(ts.parameters) + m.info[:st] = cdev(Lux.testmode(ts.states)) +end + +function fit_iter!(m, cache) + ts = cache[:train_state] + lux_loss = cache[:lux_loss] + + for d in cache[:data] + _, loss, _, ts = Training.single_train_step!( + AutoEnzyme(), lux_loss, (d[1], d[2]), ts + ) + end + + cache[:train_state] = ts + m.info[:nrounds] += 1 + return nothing +end + end \ No newline at end of file diff --git a/src/MLJ.jl b/src/MLJ.jl index 4138738..f619bb9 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,110 +3,109 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit: init, fit_iter! +import ..Fit: init, fit_iter!, _sync_params_to_model! import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema export fit, update, predict function fit( - model::LearnerTypes, - verbosity::Int, - A, - y, - w=nothing) - - Tables.istable(A) ? dtrain = DataFrame(A) : error("`A` must be a Table") - feature_names = string.(collect(Tables.schema(dtrain).names)) - @assert "_target" ∉ feature_names - dtrain._target = y - target_name = "_target" - - if !isnothing(w) - @assert "_weight" ∉ feature_names - dtrain._weight = w - weight_name = "_weight" - else - weight_name = nothing - end - offset_name = nothing - - fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) - - while fitresult.info[:nrounds] < model.nrounds - fit_iter!(fitresult, cache) - end - - report = (features=fitresult.info[:feature_names],) - return fitresult, cache, report + model::LearnerTypes, + verbosity::Int, + A, + y, + w=nothing) + + Tables.istable(A) ? dtrain = DataFrame(A) : error("`A` must be a Table") + feature_names = string.(collect(Tables.schema(dtrain).names)) + @assert "_target" ∉ feature_names + dtrain._target = y + target_name = "_target" + + if !isnothing(w) + @assert "_weight" ∉ feature_names + dtrain._weight = w + weight_name = "_weight" + else + weight_name = nothing + end + offset_name = nothing + + fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) + + while fitresult.info[:nrounds] < model.nrounds + fit_iter!(fitresult, cache) + end + _sync_params_to_model!(fitresult, cache) + + report = (features=fitresult.info[:feature_names],) + return fitresult, cache, report end function okay_to_continue(model, fitresult, cache) - return model.nrounds - fitresult.info[:nrounds] >= 0 + return model.nrounds - fitresult.info[:nrounds] >= 0 end -# For EarlyStopping.jl support MMI.iteration_parameter(::Type{<:LearnerTypes}) = :nrounds function update( - model::LearnerTypes, - verbosity::Integer, - fitresult, - cache, - A, - y, - w=nothing, + model::LearnerTypes, + verbosity::Integer, + fitresult, + cache, + A, + y, + w=nothing, ) - if okay_to_continue(model, fitresult, cache) - while fitresult.info[:nrounds] < model.nrounds - fit_iter!(fitresult, cache) + if okay_to_continue(model, fitresult, cache) + while fitresult.info[:nrounds] < model.nrounds + fit_iter!(fitresult, cache) + end + _sync_params_to_model!(fitresult, cache) + report = (features=fitresult.info[:feature_names],) + else + fitresult, cache, report = fit(model, verbosity, A, y, w) end - report = (features=fitresult.info[:feature_names],) - else - fitresult, cache, report = fit(model, verbosity, A, y, w) - end - return fitresult, cache, report + return fitresult, cache, report end -function predict(::NeuroTabRegressor, fitresult, A; device=:cpu, gpuID=0) - df = DataFrame(A) - Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df; device, gpuID) - return pred +function predict(::NeuroTabRegressor, fitresult, A) + Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") + pred = fitresult(df) + return pred end -function predict(::NeuroTabClassifier, fitresult, A; device=:cpu, gpuID=0) - df = DataFrame(A) - Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df; device, gpuID) - return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) +function predict(::NeuroTabClassifier, fitresult, A) + Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") + pred = fitresult(df) + return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end # Metadata MMI.metadata_pkg.( - (NeuroTabRegressor, NeuroTabClassifier), - name="NeuroTabModels", - uuid="1db4e0a5-a364-4b0c-897c-2bd5a4a3a1f2", - url="https://github.com/Evovest/NeuroTabModels.jl", - julia=true, - license="Apache", - is_wrapper=false, + (NeuroTabRegressor, NeuroTabClassifier), + name="NeuroTabModels", + uuid="1db4e0a5-a364-4b0c-897c-2bd5a4a3a1f2", + url="https://github.com/Evovest/NeuroTabModels.jl", + julia=true, + license="Apache", + is_wrapper=false, ) MMI.metadata_model( - NeuroTabRegressor, - input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), - target_scitype=AbstractVector{<:MMI.Continuous}, - weights=true, - path="NeuroTabModels.NeuroTabRegressor", + NeuroTabRegressor, + input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), + target_scitype=AbstractVector{<:MMI.Continuous}, + weights=true, + path="NeuroTabModels.NeuroTabRegressor", ) MMI.metadata_model( - NeuroTabClassifier, - input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), - target_scitype=AbstractVector{<:MMI.Finite}, - weights=true, - path="NeuroTabModels.NeuroTabClassifier", + NeuroTabClassifier, + input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), + target_scitype=AbstractVector{<:MMI.Finite}, + weights=true, + path="NeuroTabModels.NeuroTabClassifier", ) -end +end \ No newline at end of file diff --git a/src/data.jl b/src/data.jl index 7fdd5da..a6ee8df 100644 --- a/src/data.jl +++ b/src/data.jl @@ -7,11 +7,9 @@ import MLUtils: DataLoader using DataFrames using CategoricalArrays -using CUDA: CuIterator """ ContainerTrain - """ struct ContainerTrain{A<:AbstractMatrix,B,C,D} x::A @@ -48,7 +46,6 @@ function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B return (x, y, w, offset) end - function get_df_loader_train( df::AbstractDataFrame; feature_names, @@ -56,8 +53,7 @@ function get_df_loader_train( weight_name=nothing, offset_name=nothing, batchsize, - shuffle=true, - device=:cpu) + shuffle=true) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') @@ -83,10 +79,8 @@ function get_df_loader_train( return dtrain end - """ ContainerInfer - """ struct ContainerInfer{A<:AbstractMatrix,D} x::A @@ -114,8 +108,7 @@ function get_df_loader_infer( df::AbstractDataFrame; feature_names, offset_name=nothing, - batchsize, - device=:cpu) + batchsize) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') @@ -123,17 +116,13 @@ function get_df_loader_infer( offset = if isnothing(offset_name) nothing else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') + isa(offset_name, String) ? Float32.(df[!, offset_name]) : Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') end container = ContainerInfer(x, offset) batchsize = min(batchsize, length(container)) dinfer = DataLoader(container; shuffle=false, batchsize, partial=true, parallel=false) - if device == :gpu - return CuIterator(dinfer) - else - return dinfer - end + return dinfer end end #module \ No newline at end of file diff --git a/src/infer.jl b/src/infer.jl index e09ec39..2339035 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -4,107 +4,95 @@ using ..Data using ..Losses using ..Models -using Flux: sigmoid, softmax!, cpu, gpu, onecold +using Lux +using NNlib: sigmoid, softmax! using DataFrames: AbstractDataFrame import MLUtils: DataLoader -import CUDA: CuIterator, device! export infer """ - DL - -Union{NeuroTabModels.CuIterator, NeuroTabModels.DataLoader} -""" -const DL = Union{CuIterator,DataLoader} - -""" -infer(m::NeuroTabModel, data) + infer(m::NeuroTabModel, data::AbstractDataFrame) Return the inference of a `NeuroTabModel` over `data`, where `data` is `AbstractDataFrame`. """ -function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu, gpuID=0) - if device == :gpu - device!(gpuID) - end - m = device == :gpu ? m |> gpu : m |> cpu - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048, device) +function infer(m::NeuroTabModel, data::AbstractDataFrame) + dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) p = infer(m, dinfer) return p end +function _forward(m::NeuroTabModel, x::AbstractMatrix) + ps = m.info[:ps] + st = m.info[:st] + p, _ = Lux.apply(m.chain, x, ps, st) + if size(p, 1) == 1 + p = dropdims(p; dims=1) + end + return p +end """ (m::NeuroTabModel)(x::AbstractMatrix) (m::NeuroTabModel)(data::AbstractDataFrame) -Inference for NeuroTabModel +Inference for NeuroTabModel. """ function (m::NeuroTabModel)(x::AbstractMatrix) - p = m.chain(x) - if size(p, 1) == 1 - p = dropdims(p; dims=1) - end - return p + return _forward(m, x) end -function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu, gpuID=0) - if device == :gpu - device!(gpuID) - end - m = device == :gpu ? m |> gpu : m |> cpu - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048, device) +function (m::NeuroTabModel)(data::AbstractDataFrame) + dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) p = infer(m, dinfer) return p end - -function infer(m::NeuroTabModel{L}, data::DL) where {L<:Union{MSE,MAE}} +function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{MSE,MAE}} preds = Vector{Float32}[] for x in data - push!(preds, Vector(m(x))) + push!(preds, Vector(_forward(m, x))) end - p = vcat(preds...) - return p + return vcat(preds...) end -function infer(m::NeuroTabModel{<:LogLoss}, data::DL) +function infer(m::NeuroTabModel{<:LogLoss}, data::DataLoader) preds = Vector{Float32}[] for x in data - push!(preds, Vector(m(x))) + push!(preds, Vector(_forward(m, x))) end p = vcat(preds...) - p .= sigmoid(p) + p .= sigmoid.(p) return p end -function infer(m::NeuroTabModel{<:MLogLoss}, data::DL) +function infer(m::NeuroTabModel{<:MLogLoss}, data::DataLoader) preds = Matrix{Float32}[] for x in data - push!(preds, Matrix(m(x)')) + push!(preds, Matrix(_forward(m, x)')) end p = vcat(preds...) softmax!(p; dims=2) return p end -function infer(m::NeuroTabModel{<:GaussianMLE}, data::DL) +function infer(m::NeuroTabModel{<:GaussianMLE}, data::DataLoader) preds = Matrix{Float32}[] for x in data - push!(preds, Matrix(m(x)')) + push!(preds, Matrix(_forward(m, x)')) end p = vcat(preds...) - p[:, 2] .= exp.(p[:, 2]) # reproject log(σ) into σ + p[:, 2] .= exp.(p[:, 2]) return p end -function infer(m::NeuroTabModel{L}, data::DL) where {L<:Union{Tweedie}} +function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{Tweedie}} preds = Vector{Float32}[] for x in data - push!(preds, Vector(m(x))) + push!(preds, Vector(_forward(m, x))) end p = vcat(preds...) p .= exp.(p) return p end -end # module +end # module \ No newline at end of file diff --git a/src/losses.jl b/src/losses.jl index bec273e..3cbd4bf 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -4,7 +4,8 @@ export get_loss_fn, get_loss_type export LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE, Tweedie import Statistics: mean, std -import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid, onehotbatch +import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid +import OneHotArrays: onehotbatch abstract type LossType end abstract type MSE <: LossType end @@ -36,15 +37,15 @@ end function logloss(m, x, y) p = m(x) - mean((1 .- y) .* p .- logσ.(p)) + mean((1 .- y) .* p .- logsigmoid.(p)) end function logloss(m, x, y, w) p = m(x) - sum(w .* ((1 .- y) .* p .- logσ.(p))) / sum(w) + sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end function logloss(m, x, y, w, offset) p = m(x) .+ offset - sum(w .* ((1 .- y) .* p .- logσ.(p))) / sum(w) + sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end function tweedie(m, x, y) @@ -72,17 +73,17 @@ end function mlogloss(m, x, y) p = logsoftmax(m(x); dims=1) k = size(p, 1) - mean(-sum(onehotbatch(y, 1:k) .* p; dims=1)) + mean(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1)) end function mlogloss(m, x, y, w) p = logsoftmax(m(x); dims=1) k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) + sum(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1) .* w) / sum(w) end function mlogloss(m, x, y, w, offset) p = logsoftmax(m(x) .+ offset; dims=1) k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) + sum(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1) .* w) / sum(w) end gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}) where {T} = diff --git a/src/metrics.jl b/src/metrics.jl index 3110010..6955f10 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -3,7 +3,8 @@ module Metrics export metric_dict, is_maximise import Statistics: mean, std -import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid, onehotbatch +import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid +import OneHotArrays: onehotbatch """ mse(x, y; agg=mean) @@ -49,17 +50,17 @@ end """ function logloss(m, x, y; agg=mean) p = m(x) - metric = agg((1 .- y) .* p .- logσ.(p)) + metric = agg((1 .- y) .* p .- logsigmoid.(p)) return metric end function logloss(m, x, y, w; agg=mean) p = m(x) - metric = agg(((1 .- y) .* p .- logσ.(p)) .* w) + metric = agg(((1 .- y) .* p .- logsigmoid.(p)) .* w) return metric end function logloss(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - metric = agg(((1 .- y) .* p .- logσ.(p)) .* w) + metric = agg(((1 .- y) .* p .- logsigmoid.(p)) .* w) return metric end @@ -100,26 +101,25 @@ end function mlogloss(m, x, y; agg=mean) p = logsoftmax(m(x); dims=1) k = size(p, 1) - raw = dropdims(-sum(onehotbatch(y, 1:k) .* p; dims=1); dims=1) + raw = dropdims(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1); dims=1) metric = agg(raw) return metric end function mlogloss(m, x, y, w; agg=mean) p = logsoftmax(m(x); dims=1) k = size(p, 1) - raw = dropdims(-sum(onehotbatch(y, 1:k) .* p; dims=1); dims=1) + raw = dropdims(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1); dims=1) metric = agg(raw .* w) return metric end function mlogloss(m, x, y, w, offset; agg=mean) p = logsoftmax(m(x) .+ offset; dims=1) k = size(p, 1) - raw = dropdims(-sum(onehotbatch(y, 1:k) .* p; dims=1); dims=1) + raw = dropdims(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1); dims=1) metric = agg(raw .* w) return metric end - """ gaussian_mle(μ::T, σ::T, y::T, w::T) where {T<:AbstractFloat} """ From 88cf9a70ebac6eeb956a34ea4be54f1087f01819 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 16 Feb 2026 08:54:34 +0530 Subject: [PATCH 042/120] simplify get_device Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 23dda8e..9ffd43a 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -24,20 +24,9 @@ include("callback.jl") using .CallBacks function _get_device(config) - device = Symbol(config.device) - if device == :gpu - try - Reactant.set_default_backend("gpu") - return reactant_device() - catch - @warn "GPU not available, falling back to CPU" - Reactant.set_default_backend("cpu") - return reactant_device() - end - else - Reactant.set_default_backend("cpu") - return reactant_device() - end + backend = config.device == :gpu ? "gpu" : "cpu" + Reactant.set_default_backend(backend) + return reactant_device() end function _get_lux_loss(L) From 24f3fdeb2869953ec368d3f14572bfe29f65d953 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Mon, 16 Feb 2026 00:37:58 -0500 Subject: [PATCH 043/120] up --- Project.toml | 2 ++ benchmarks/benchmark_mse.jl | 5 +++-- src/Fit/fit.jl | 22 +++++++++++++++------- src/MLJ.jl | 6 +++--- 4 files changed, 23 insertions(+), 12 deletions(-) diff --git a/Project.toml b/Project.toml index 5e15b77..6265a1a 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -22,6 +23,7 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 8988d2a..c1f88d9 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -43,9 +43,10 @@ learner = NeuroTabRegressor( device=:gpu ) -# desktop gpu: 13.476383 seconds (26.42 M allocations: 5.990 GiB, 9.44% gc time) +# Reactant GPU: 5.970480 seconds (2.33 M allocations: 5.242 GiB, 3.80% gc time, 0.00% compilation time) +# Zygote GPU: 9.855853 seconds (27.92 M allocations: 6.005 GiB, 3.58% gc time) # 13.557744 seconds (26.40 M allocations: 5.989 GiB, 9.60% gc time) -@time m = NeuroTabModels.fit( +@time m, ts = NeuroTabModels.fit( learner, dtrain; target_name, diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 30fd06e..728da24 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -16,6 +16,8 @@ import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, using Lux using Enzyme, Reactant +# using Zygote +# using Mooncake using DataFrames using CategoricalArrays @@ -35,8 +37,10 @@ function init( device = config.device batchsize = config.batchsize nfeats = length(feature_names) - loss = get_loss_fn(config.loss) L = get_loss_type(config.loss) + # loss = get_loss_fn(config.loss) + loss = MSELoss() + # loss = BinaryCrossEntropyLoss(; logits=Val(true)), target_levels = nothing target_isordered = false @@ -52,8 +56,12 @@ function init( Reactant.set_default_backend("gpu") dev = reactant_device() + backend = AutoReactant() # dev = cpu_device() # dev = gpu_device() + # backend = AutoZygote() + # backend = AutoMooncake() + data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev info = Dict( @@ -69,8 +77,9 @@ function init( rng = Xoshiro(config.seed) ps, st = Lux.setup(rng, m.chain) |> dev opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) + ts = Training.TrainState(m.chain, ps, st, opt) - cache = (data=data, ps=ps, st=st, opt=opt, info=info) + cache = (data=data, ts=ts, loss=loss, backend=backend, info=info) return m, cache end @@ -138,13 +147,12 @@ function fit( (verbosity > 0) && @info "Init training" end - ts = Training.TrainState(m.chain, cache[:ps], cache[:st], cache[:opt]) + ts = cache[:ts] while m.info[:nrounds] < config.nrounds for d in cache[:data] gs, loss, stats, ts = Training.single_train_step!( - AutoEnzyme(), - MSELoss(), - # BinaryCrossEntropyLoss(; logits=Val(true)), + cache.backend, + cache.loss, (d[1], d[2]), ts ) @@ -163,7 +171,7 @@ function fit( end end m.info[:logger] = logger - return m + return m, ts end end \ No newline at end of file diff --git a/src/MLJ.jl b/src/MLJ.jl index 4138738..9daa549 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,7 +3,7 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit: init, fit_iter! +import ..Fit: init import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema @@ -34,7 +34,7 @@ function fit( fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) while fitresult.info[:nrounds] < model.nrounds - fit_iter!(fitresult, cache) + # fit_iter!(fitresult, cache) end report = (features=fitresult.info[:feature_names],) @@ -59,7 +59,7 @@ function update( ) if okay_to_continue(model, fitresult, cache) while fitresult.info[:nrounds] < model.nrounds - fit_iter!(fitresult, cache) + # fit_iter!(fitresult, cache) end report = (features=fitresult.info[:feature_names],) else From be881a3dfd4e18e41bf53db6d3b4da4ef3a70791 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 16 Feb 2026 12:31:28 +0530 Subject: [PATCH 044/120] model cleanup Signed-off-by: AdityaPandeyCN --- Project.toml | 2 + src/MLJ.jl | 22 +++++---- src/models/NeuroTree/model.jl | 79 +++++------------------------- src/models/NeuroTree/neurotrees.jl | 57 ++++----------------- 4 files changed, 37 insertions(+), 123 deletions(-) diff --git a/Project.toml b/Project.toml index 23f4c6b..96f67c3 100644 --- a/Project.toml +++ b/Project.toml @@ -9,11 +9,13 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" diff --git a/src/MLJ.jl b/src/MLJ.jl index f619bb9..5162ee4 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,7 +3,8 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit: init, fit_iter!, _sync_params_to_model! +import ..Fit +import ..Fit: init, fit_iter!, _sync_to_cpu!, _run_compiled_loop! import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema @@ -33,10 +34,9 @@ function fit( fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) - while fitresult.info[:nrounds] < model.nrounds - fit_iter!(fitresult, cache) - end - _sync_params_to_model!(fitresult, cache) + # Use the optimized chunked loop + _run_compiled_loop!(fitresult, cache, model.nrounds) + _sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) return fitresult, cache, report @@ -58,10 +58,14 @@ function update( w=nothing, ) if okay_to_continue(model, fitresult, cache) - while fitresult.info[:nrounds] < model.nrounds - fit_iter!(fitresult, cache) + rounds_to_add = model.nrounds - fitresult.info[:nrounds] + + if rounds_to_add > 0 + # Optimized update using chunked loop + _run_compiled_loop!(fitresult, cache, rounds_to_add) end - _sync_params_to_model!(fitresult, cache) + + _sync_to_cpu!(fitresult, cache) report = (features=fitresult.info[:feature_names],) else fitresult, cache, report = fit(model, verbosity, A, y, w) @@ -81,7 +85,7 @@ function predict(::NeuroTabClassifier, fitresult, A) return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end -# Metadata +# Metadata (Unchanged) MMI.metadata_pkg.( (NeuroTabRegressor, NeuroTabClassifier), name="NeuroTabModels", diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 23b3fa2..f78a54f 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -22,40 +22,36 @@ function NeuroTree((feats, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, Float32(init_scale)) end -# Define the Lux interface function LuxCore.initialparameters(rng::AbstractRNG, l::NeuroTree) return ( - w=Float32.((rand(l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # w - b=zeros(Float32, l.nodes * l.trees), # b - s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), # s - p=Float32.(randn(l.outs, l.leaves * l.trees) .* l.init_scale), # p + w=Float32.((rand(rng, l.nodes * l.trees, l.feats) .- 0.5) ./ 4), + b=zeros(Float32, l.nodes * l.trees), + s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), + p=Float32.(randn(rng, l.outs, l.leaves * l.trees) .* l.init_scale), ) end function LuxCore.initialstates(rng::AbstractRNG, l::NeuroTree) return ( - ml=get_logits_mask(Val(l.tree_type), l.depth), - ms=get_softplus_mask(Val(l.tree_type), l.depth) + ml=Float32.(get_logits_mask(Val(l.tree_type), l.depth)), + ms=Float32.(get_softplus_mask(Val(l.tree_type), l.depth)), ) end function (l::NeuroTree)(x, ps, st) if l.scaler - nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] + nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) else - nw = (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] + nw = (l.actA(ps.w) * x .+ ps.b) end - nw = reshape(nw, size(st.ml, 2), :) # [NT,B] => [N,TB] - lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) # [N,TB] => [L,TB] - lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] - y = ps.p * lw ./ l.trees # [P,LT] * [LT,B] => [P,B] + nw = reshape(nw, size(st.ml, 2), :) + lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) + lw = reshape(lw, :, size(x, 2)) + y = ps.p * lw ./ l.trees return y, st end -""" - get_logits_mask(::Val{:binary}, depth::Integer) -""" function get_logits_mask(::Val{:binary}, depth::Integer) nodes = 2^depth - 1 leaves = 2^depth @@ -84,9 +80,6 @@ function get_logits_mask(::Val{:oblivious}, depth::Integer) return mask end -""" - get_softplus_mask(::Val{:binary}, depth::Integer) -""" function get_softplus_mask(::Val{:binary}, depth::Integer) nodes = 2^depth - 1 leaves = 2^depth @@ -105,50 +98,4 @@ function get_softplus_mask(::Val{:oblivious}, depth::Integer) leaves = 2^depth mask = ones(Bool, leaves, depth) return mask -end - - -""" - StackTree -A StackTree is made of a collection of NeuroTree. -""" -struct StackTree - trees::Vector{NeuroTree} -end - -function StackTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, stack_size=1, hidden_size=8, actA=identity, scaler=true, init_scale=1e-1) - @assert stack_size == 1 || hidden_size >= outs - trees = [] - for i in 1:stack_size - if i == 1 - if i < stack_size - tree = NeuroTree(ins => hidden_size; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - else - tree = NeuroTree(ins => outs; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - end - elseif i < stack_size - tree = NeuroTree(hidden_size => hidden_size; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - else - tree = NeuroTree(hidden_size => outs; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - end - end - m = StackTree(trees) - return m -end - -function (m::StackTree)(x::AbstractMatrix) - p = m.trees[1](x) - for i in 2:length(m.trees) - if i < length(m.trees) - p = p .+ m.trees[i](p) - else - _p = m.trees[i](p) - p = view(p, 1:size(_p, 1), :) .+ _p - end - end - return p -end +end \ No newline at end of file diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 8db3ce1..f567cd7 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -3,16 +3,9 @@ module NeuroTrees export NeuroTreeConfig using Random -import .Threads: @threads - -# using CUDA -# import Flux -# import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu -# import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch, glorot_uniform -# import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel - using Lux -using NNlib: softplus, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh +using LuxCore +using NNlib: softplus, sigmoid, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture @@ -33,8 +26,6 @@ struct NeuroTreeConfig <: Architecture end function NeuroTreeConfig(; kwargs...) - - # defaults arguments args = Dict{Symbol,Any}( :tree_type => :binary, :actA => :identity, @@ -80,49 +71,29 @@ function NeuroTreeConfig(; kwargs...) end function (config::NeuroTreeConfig)(; nfeats, outsize) - if config.MLE_tree_split && outsize == 2 outsize ÷= 2 chain = Chain( BatchNorm(nfeats), Parallel( vcat, - StackTree(nfeats => outsize; + NeuroTree(nfeats => outsize; tree_type=config.tree_type, depth=config.depth, - ntrees=config.ntrees, - proj_size=config.proj_size, - stack_size=config.stack_size, - hidden_size=config.hidden_size, + trees=config.ntrees, actA=act_dict[config.actA], scaler=config.scaler, init_scale=config.init_scale), - StackTree(nfeats => outsize; + NeuroTree(nfeats => outsize; tree_type=config.tree_type, depth=config.depth, - ntrees=config.ntrees, - proj_size=config.proj_size, - stack_size=config.stack_size, - hidden_size=config.hidden_size, + trees=config.ntrees, actA=act_dict[config.actA], scaler=config.scaler, - init_scale=config.init_scale) + init_scale=config.init_scale), ) ) else - # chain = Chain( - # BatchNorm(nfeats), - # StackTree(nfeats => outsize; - # tree_type=config.tree_type, - # depth=config.depth, - # ntrees=config.ntrees, - # proj_size=config.proj_size, - # stack_size=config.stack_size, - # hidden_size=config.hidden_size, - # actA=act_dict[config.actA], - # scaler=config.scaler, - # init_scale=config.init_scale) - # ) chain = Chain( BatchNorm(nfeats), NeuroTree(nfeats => outsize; @@ -136,28 +107,18 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) end end - function _identity_act(x) return x ./ sum(abs.(x), dims=2) end function _tanh_act(x) - x = Flux.tanh_fast.(x) + x = tanh_fast.(x) return x ./ sum(abs.(x), dims=2) end function _hardtanh_act(x) - x = Flux.hardtanh.(x) + x = hardtanh.(x) return x ./ sum(abs.(x), dims=2) end -""" - act_dict = Dict( - :identity => _identity_act, - :tanh => _tanh_act, - :hardtanh => _hardtanh_act, - ) - -Dictionary mapping features activation name to their function. -""" const act_dict = Dict( :identity => _identity_act, :tanh => _tanh_act, From 9dad435dee0b8b0967d770b99dcf69d3d051889a Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 16 Feb 2026 13:01:38 +0530 Subject: [PATCH 045/120] fix imports Signed-off-by: AdityaPandeyCN --- src/MLJ.jl | 151 ++++++++++++++++++++++++++--------------------------- 1 file changed, 75 insertions(+), 76 deletions(-) diff --git a/src/MLJ.jl b/src/MLJ.jl index 5162ee4..6cce9be 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,113 +3,112 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit -import ..Fit: init, fit_iter!, _sync_to_cpu!, _run_compiled_loop! +import ..Fit: init, fit_iter!, _sync_params_to_model! import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema export fit, update, predict function fit( - model::LearnerTypes, - verbosity::Int, - A, - y, - w=nothing) - - Tables.istable(A) ? dtrain = DataFrame(A) : error("`A` must be a Table") - feature_names = string.(collect(Tables.schema(dtrain).names)) - @assert "_target" ∉ feature_names - dtrain._target = y - target_name = "_target" - - if !isnothing(w) - @assert "_weight" ∉ feature_names - dtrain._weight = w - weight_name = "_weight" - else - weight_name = nothing - end - offset_name = nothing - - fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) - - # Use the optimized chunked loop - _run_compiled_loop!(fitresult, cache, model.nrounds) - _sync_to_cpu!(fitresult, cache) - - report = (features=fitresult.info[:feature_names],) - return fitresult, cache, report + model::LearnerTypes, + verbosity::Int, + A, + y, + w=nothing) + + Tables.istable(A) ? dtrain = DataFrame(A) : error("`A` must be a Table") + feature_names = string.(collect(Tables.schema(dtrain).names)) + @assert "_target" ∉ feature_names + dtrain._target = y + target_name = "_target" + + if !isnothing(w) + @assert "_weight" ∉ feature_names + dtrain._weight = w + weight_name = "_weight" + else + weight_name = nothing + end + offset_name = nothing + + fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) + + while fitresult.info[:nrounds] < model.nrounds + fit_iter!(fitresult, cache) + end + + _sync_params_to_model!(fitresult, cache) + report = (features=fitresult.info[:feature_names],) + return fitresult, cache, report end function okay_to_continue(model, fitresult, cache) - return model.nrounds - fitresult.info[:nrounds] >= 0 + return model.nrounds - fitresult.info[:nrounds] >= 0 end +# For EarlyStopping.jl support MMI.iteration_parameter(::Type{<:LearnerTypes}) = :nrounds function update( - model::LearnerTypes, - verbosity::Integer, - fitresult, - cache, - A, - y, - w=nothing, + model::LearnerTypes, + verbosity::Integer, + fitresult, + cache, + A, + y, + w=nothing, ) - if okay_to_continue(model, fitresult, cache) - rounds_to_add = model.nrounds - fitresult.info[:nrounds] - - if rounds_to_add > 0 - # Optimized update using chunked loop - _run_compiled_loop!(fitresult, cache, rounds_to_add) - end - - _sync_to_cpu!(fitresult, cache) - report = (features=fitresult.info[:feature_names],) - else - fitresult, cache, report = fit(model, verbosity, A, y, w) + if okay_to_continue(model, fitresult, cache) + while fitresult.info[:nrounds] < model.nrounds + fit_iter!(fitresult, cache) end - return fitresult, cache, report + _sync_params_to_model!(fitresult, cache) + report = (features=fitresult.info[:feature_names],) + else + fitresult, cache, report = fit(model, verbosity, A, y, w) + end + return fitresult, cache, report end function predict(::NeuroTabRegressor, fitresult, A) - Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df) - return pred + df = DataFrame(A) + Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") + pred = fitresult(df) + return pred end function predict(::NeuroTabClassifier, fitresult, A) - Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df) - return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) + df = DataFrame(A) + Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") + pred = fitresult(df) + return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end -# Metadata (Unchanged) +# Metadata MMI.metadata_pkg.( - (NeuroTabRegressor, NeuroTabClassifier), - name="NeuroTabModels", - uuid="1db4e0a5-a364-4b0c-897c-2bd5a4a3a1f2", - url="https://github.com/Evovest/NeuroTabModels.jl", - julia=true, - license="Apache", - is_wrapper=false, + (NeuroTabRegressor, NeuroTabClassifier), + name="NeuroTabModels", + uuid="1db4e0a5-a364-4b0c-897c-2bd5a4a3a1f2", + url="https://github.com/Evovest/NeuroTabModels.jl", + julia=true, + license="Apache", + is_wrapper=false, ) MMI.metadata_model( - NeuroTabRegressor, - input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), - target_scitype=AbstractVector{<:MMI.Continuous}, - weights=true, - path="NeuroTabModels.NeuroTabRegressor", + NeuroTabRegressor, + input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), + target_scitype=AbstractVector{<:MMI.Continuous}, + weights=true, + path="NeuroTabModels.NeuroTabRegressor", ) MMI.metadata_model( - NeuroTabClassifier, - input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), - target_scitype=AbstractVector{<:MMI.Finite}, - weights=true, - path="NeuroTabModels.NeuroTabClassifier", + NeuroTabClassifier, + input_scitype=MMI.Table(MMI.Continuous, MMI.Count, MMI.OrderedFactor), + target_scitype=AbstractVector{<:MMI.Finite}, + weights=true, + path="NeuroTabModels.NeuroTabClassifier", ) end \ No newline at end of file From b6b47f4fa5b06dd3bc7510983563a6b8497f5f1e Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 16 Feb 2026 13:18:21 +0530 Subject: [PATCH 046/120] fix imports Signed-off-by: AdityaPandeyCN --- src/MLJ.jl | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/MLJ.jl b/src/MLJ.jl index 6cce9be..6a45c4b 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -71,14 +71,12 @@ function update( end function predict(::NeuroTabRegressor, fitresult, A) - df = DataFrame(A) Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") pred = fitresult(df) return pred end function predict(::NeuroTabClassifier, fitresult, A) - df = DataFrame(A) Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") pred = fitresult(df) return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) From 674ef555d5e50fdd8cbe947b3d769199aebf87bd Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 16 Feb 2026 19:47:07 +0530 Subject: [PATCH 047/120] add gpu inference support Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 4 +- src/infer.jl | 143 +++++++++++++++++++++++------------------ src/metrics.jl | 171 ++++++++++++++++++++++++------------------------- test/core.jl | 41 ++++++++++-- 4 files changed, 204 insertions(+), 155 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 9ffd43a..4206c81 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -88,7 +88,8 @@ function init( :nrounds => 0, :feature_names => feature_names, :target_levels => target_levels, - :target_isordered => target_isordered) + :target_isordered => target_isordered, + :device => config.device) chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) @@ -155,6 +156,7 @@ function fit( m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) + # initialize callback and logger if tracking eval data logger = nothing if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) diff --git a/src/infer.jl b/src/infer.jl index 2339035..1cd862c 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -5,94 +5,113 @@ using ..Losses using ..Models using Lux +using Lux: cpu_device, reactant_device +using Reactant using NNlib: sigmoid, softmax! using DataFrames: AbstractDataFrame import MLUtils: DataLoader export infer -""" - infer(m::NeuroTabModel, data::AbstractDataFrame) - -Return the inference of a `NeuroTabModel` over `data`, where `data` is `AbstractDataFrame`. -""" -function infer(m::NeuroTabModel, data::AbstractDataFrame) - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) - p = infer(m, dinfer) - return p -end - -function _forward(m::NeuroTabModel, x::AbstractMatrix) - ps = m.info[:ps] - st = m.info[:st] - p, _ = Lux.apply(m.chain, x, ps, st) - if size(p, 1) == 1 - p = dropdims(p; dims=1) +function _get_device(device::Symbol) + if device == :gpu + Reactant.set_default_backend("gpu") + return reactant_device() + else + return cpu_device() end - return p end """ - (m::NeuroTabModel)(x::AbstractMatrix) - (m::NeuroTabModel)(data::AbstractDataFrame) + _setup_infer(m, data, device) -> (ps, st, model_fn, batches) -Inference for NeuroTabModel. """ -function (m::NeuroTabModel)(x::AbstractMatrix) - return _forward(m, x) -end -function (m::NeuroTabModel)(data::AbstractDataFrame) - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) - p = infer(m, dinfer) - return p -end +function _setup_infer(m::NeuroTabModel, data, device::Symbol) + dev = _get_device(device) + ps = dev(m.info[:ps]) + st = dev(m.info[:st]) + batches = collect(data) |> dev -function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{MSE,MAE}} - preds = Vector{Float32}[] - for x in data - push!(preds, Vector(_forward(m, x))) + if device == :gpu + compiled_model = Reactant.@compile m.chain(first(batches), ps, st) + return ps, st, compiled_model, batches + else + return ps, st, m.chain, batches end - return vcat(preds...) end -function infer(m::NeuroTabModel{<:LogLoss}, data::DataLoader) - preds = Vector{Float32}[] - for x in data - push!(preds, Vector(_forward(m, x))) - end - p = vcat(preds...) - p .= sigmoid.(p) - return p +# Run single forward pass on device, pull result back to CPU +function _forward_batch(model_fn, x, ps, st) + y, _ = model_fn(x, ps, st) + return cpu_device()(y) end -function infer(m::NeuroTabModel{<:MLogLoss}, data::DataLoader) - preds = Matrix{Float32}[] - for x in data - push!(preds, Matrix(_forward(m, x)')) - end - p = vcat(preds...) +function _postprocess(::Type{<:Union{MSE,MAE}}, raw_preds::Vector) + return vcat([vec(p) for p in raw_preds]...) +end +function _postprocess(::Type{<:LogLoss}, raw_preds::Vector) + p = vcat([vec(p) for p in raw_preds]...) + return sigmoid.(p) +end +function _postprocess(::Type{<:MLogLoss}, raw_preds::Vector) + # Transpose from [class, batch] -> [batch, class], then softmax + p = vcat([Matrix(p') for p in raw_preds]...) softmax!(p; dims=2) return p end - -function infer(m::NeuroTabModel{<:GaussianMLE}, data::DataLoader) - preds = Matrix{Float32}[] - for x in data - push!(preds, Matrix(_forward(m, x)')) - end - p = vcat(preds...) +function _postprocess(::Type{<:GaussianMLE}, raw_preds::Vector) + # Transpose, then exponentiate the σ column (log(σ) -> σ) + p = vcat([Matrix(p') for p in raw_preds]...) p[:, 2] .= exp.(p[:, 2]) return p end +function _postprocess(::Type{<:Tweedie}, raw_preds::Vector) + p = vcat([vec(p) for p in raw_preds]...) + return exp.(p) +end -function infer(m::NeuroTabModel{L}, data::DataLoader) where {L<:Union{Tweedie}} - preds = Vector{Float32}[] - for x in data - push!(preds, Vector(_forward(m, x))) +""" + infer(m::NeuroTabModel, data::DataLoader; device=:cpu) + +""" +function infer(m::NeuroTabModel{L}, data::DataLoader; device=:cpu) where {L} + ps, st, model_fn, batches = _setup_infer(m, data, device) + raw_preds = [_forward_batch(model_fn, x, ps, st) for x in batches] + return _postprocess(L, raw_preds) +end + +""" + infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu) + +""" +function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu) + dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) + return infer(m, dinfer; device=device) +end + +""" + (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) + +""" +function (m::NeuroTabModel{L})(x::AbstractMatrix; device=:cpu) where {L} + dev = _get_device(device) + ps, st = dev(m.info[:ps]), dev(m.info[:st]) + x_dev = dev(x) + if device == :gpu + model_fn = Reactant.@compile m.chain(x_dev, ps, st) + else + model_fn = m.chain end - p = vcat(preds...) - p .= exp.(p) - return p + raw_pred = _forward_batch(model_fn, x_dev, ps, st) + return _postprocess(L, [raw_pred]) +end + +""" + (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) + +""" +function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) + return infer(m, data; device=device) end end # module \ No newline at end of file diff --git a/src/metrics.jl b/src/metrics.jl index 6955f10..2efa0e5 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -1,153 +1,148 @@ module Metrics -export metric_dict, is_maximise +export metric_dict, is_maximise, get_metric import Statistics: mean, std import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid import OneHotArrays: onehotbatch """ - mse(x, y; agg=mean) - mse(x, y, w; agg=mean) - mse(x, y, w, offset; agg=mean) + mse(m, x, y; agg=mean) + mse(m, x, y, w; agg=mean) + mse(m, x, y, w, offset; agg=mean) """ function mse(m, x, y; agg=mean) - metric = agg((m(x) .- y) .^ 2) - return metric + return agg((vec(m(x)) .- vec(y)) .^ 2) end function mse(m, x, y, w; agg=mean) - metric = agg((m(x) .- y) .^ 2 .* w) - return metric + return agg((vec(m(x)) .- vec(y)) .^ 2 .* vec(w)) end function mse(m, x, y, w, offset; agg=mean) - metric = agg((m(x) .+ offset .- y) .^ 2 .* w) - return metric + return agg((vec(m(x)) .+ vec(offset) .- vec(y)) .^ 2 .* vec(w)) end """ - mae(x, y; agg=mean) - mae(x, y, w; agg=mean) - mae(x, y, w, offset; agg=mean) + mae(m, x, y; agg=mean) + mae(m, x, y, w; agg=mean) + mae(m, x, y, w, offset; agg=mean) """ function mae(m, x, y; agg=mean) - metric = agg(abs.(m(x) .- y)) - return metric + return agg(abs.(vec(m(x)) .- vec(y))) end function mae(m, x, y, w; agg=mean) - metric = agg(abs.(m(x) .- y) .* w) - return metric + return agg(abs.(vec(m(x)) .- vec(y)) .* vec(w)) end function mae(m, x, y, w, offset; agg=mean) - metric = agg(abs.(m(x) .+ offset .- y) .* w) - return metric + return agg(abs.(vec(m(x)) .+ vec(offset) .- vec(y)) .* vec(w)) end - """ - logloss(x, y; agg=mean) - logloss(x, y, w; agg=mean) - logloss(x, y, w, offset; agg=mean) + logloss(m, x, y; agg=mean) + logloss(m, x, y, w; agg=mean) + logloss(m, x, y, w, offset; agg=mean) """ function logloss(m, x, y; agg=mean) - p = m(x) - metric = agg((1 .- y) .* p .- logsigmoid.(p)) - return metric + p = vec(m(x)) + y = vec(y) + return agg((1 .- y) .* p .- logsigmoid.(p)) end function logloss(m, x, y, w; agg=mean) - p = m(x) - metric = agg(((1 .- y) .* p .- logsigmoid.(p)) .* w) - return metric + p = vec(m(x)) + y = vec(y) + return agg(((1 .- y) .* p .- logsigmoid.(p)) .* vec(w)) end function logloss(m, x, y, w, offset; agg=mean) - p = m(x) .+ offset - metric = agg(((1 .- y) .* p .- logsigmoid.(p)) .* w) - return metric + p = vec(m(x)) .+ vec(offset) + y = vec(y) + return agg(((1 .- y) .* p .- logsigmoid.(p)) .* vec(w)) end - """ - tweedie(x, y; agg=mean) - tweedie(x, y, w; agg=mean) - tweedie(x, y, w, offset; agg=mean) + tweedie(m, x, y; agg=mean) + tweedie(m, x, y, w; agg=mean) + tweedie(m, x, y, w, offset; agg=mean) """ function tweedie(m, x, y; agg=mean) rho = eltype(x)(1.5) - p = exp.(m(x)) - agg(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) + p = exp.(vec(m(x))) + y = vec(y) + return agg(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + + p .^ (2 - rho) / (2 - rho))) end -function tweedie(m, x, y, w) - agg = mean +function tweedie(m, x, y, w; agg=mean) rho = eltype(x)(1.5) - p = exp.(m(x)) - agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) + p = exp.(vec(m(x))) + y = vec(y) + w = vec(w) + return agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + + p .^ (2 - rho) / (2 - rho))) end function tweedie(m, x, y, w, offset; agg=mean) rho = eltype(x)(1.5) - p = exp.(m(x) .+ offset) - agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) + p = exp.(vec(m(x)) .+ vec(offset)) + y = vec(y) + w = vec(w) + return agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + + p .^ (2 - rho) / (2 - rho))) end """ - mlogloss(x, y; agg=mean) - mlogloss(x, y, w; agg=mean) - mlogloss(x, y, w, offset; agg=mean) + mlogloss(m, x, y; agg=mean) + mlogloss(m, x, y, w; agg=mean) + mlogloss(m, x, y, w, offset; agg=mean) """ function mlogloss(m, x, y; agg=mean) - p = logsoftmax(m(x); dims=1) - k = size(p, 1) - raw = dropdims(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1); dims=1) - metric = agg(raw) - return metric + p = m(x) + p_t = permutedims(p, (2, 1)) + lsm = logsoftmax(p_t; dims=1) + k = size(lsm, 1) + y_oh = onehotbatch(vec(y), UInt32(1):UInt32(k)) + raw = -sum(y_oh .* lsm; dims=1) + return agg(vec(raw)) end function mlogloss(m, x, y, w; agg=mean) - p = logsoftmax(m(x); dims=1) - k = size(p, 1) - raw = dropdims(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1); dims=1) - metric = agg(raw .* w) - return metric + p = m(x) + p_t = permutedims(p, (2, 1)) + lsm = logsoftmax(p_t; dims=1) + k = size(lsm, 1) + y_oh = onehotbatch(vec(y), UInt32(1):UInt32(k)) + raw = -sum(y_oh .* lsm; dims=1) + return agg(vec(raw) .* vec(w)) end function mlogloss(m, x, y, w, offset; agg=mean) - p = logsoftmax(m(x) .+ offset; dims=1) - k = size(p, 1) - raw = dropdims(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1); dims=1) - metric = agg(raw .* w) - return metric + p = m(x) .+ offset + p_t = permutedims(p, (2, 1)) + lsm = logsoftmax(p_t; dims=1) + k = size(lsm, 1) + y_oh = onehotbatch(vec(y), UInt32(1):UInt32(k)) + raw = -sum(y_oh .* lsm; dims=1) + return agg(vec(raw) .* vec(w)) end -""" - gaussian_mle(μ::T, σ::T, y::T, w::T) where {T<:AbstractFloat} -""" -gaussian_mle(μ::T, σ::T, y::T) where {T<:AbstractFloat} = - -σ - (y - μ)^2 / (2 * max(T(2.0f-7), exp(2 * σ))) - -gaussian_mle(μ::T, σ::T, y::T, w::T) where {T<:AbstractFloat} = - (-σ - (y - μ)^2 / (2 * max(T(2.0f-7), exp(2 * σ)))) * w +gaussian_loss_elt(μ, σ, y) = -σ - (y - μ)^2 / (2 * max(2.0f-7, exp(2 * σ))) -"""" - gaussian_mle(x, y; agg=mean) - gaussian_mle(x, y, w; agg=mean) - gaussian_mle(x, y, w, offset; agg=mean) +""" + gaussian_mle(m, x, y; agg=mean) + gaussian_mle(m, x, y, w; agg=mean) + gaussian_mle(m, x, y, w, offset; agg=mean) """ function gaussian_mle(m, x, y; agg=mean) p = m(x) - metric = agg(gaussian_mle.(view(p, 1, :), view(p, 2, :), y)) - return metric + μ = view(p, :, 1) + σ = view(p, :, 2) + return agg(gaussian_loss_elt.(μ, σ, vec(y))) end function gaussian_mle(m, x, y, w; agg=mean) p = m(x) - metric = agg(gaussian_mle.(view(p, 1, :), view(p, 2, :), y, w)) - return metric + μ = view(p, :, 1) + σ = view(p, :, 2) + return agg(gaussian_loss_elt.(μ, σ, vec(y)) .* vec(w)) end function gaussian_mle(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - metric = agg(gaussian_mle.(view(p, 1, :), view(p, 2, :), y, w)) - return metric + μ = view(p, :, 1) + σ = view(p, :, 2) + return agg(gaussian_loss_elt.(μ, σ, vec(y)) .* vec(w)) end function get_metric(m, f::Function, data) @@ -158,7 +153,7 @@ function get_metric(m, f::Function, data) if length(d) >= 3 ws += sum(d[3]) else - ws += last(size(d[2])) + ws += size(d[2], ndims(d[2])) end end metric = metric / ws @@ -181,4 +176,4 @@ is_maximise(::typeof(mlogloss)) = false is_maximise(::typeof(gaussian_mle)) = true is_maximise(::typeof(tweedie)) = false -end +end \ No newline at end of file diff --git a/test/core.jl b/test/core.jl index 096f250..b14e774 100644 --- a/test/core.jl +++ b/test/core.jl @@ -1,5 +1,5 @@ @testset "Core - data iterators" begin - + end @testset "Core - internals test" begin @@ -36,7 +36,6 @@ end ) m = NeuroTabModel(L, chain, info) - end @testset "Core - Regression" begin @@ -64,6 +63,7 @@ end lr=1e-1, ) + # Test without eval data m = NeuroTabModels.fit( learner, dtrain; @@ -71,6 +71,13 @@ end feature_names ) + # Test inference output + p = m(dtrain) + @test p isa Vector{Float32} + @test length(p) == nrow(dtrain) + @test all(isfinite, p) + + # Test with eval data and early stopping m = NeuroTabModels.fit( learner, dtrain; @@ -79,6 +86,17 @@ end deval, ) + # Test eval metric was tracked + @test haskey(m.info, :logger) + @test !isnothing(m.info[:logger]) + @test length(m.info[:logger][:metrics][:metric]) > 0 + + # Test inference on eval data + peval = m(deval) + @test peval isa Vector{Float32} + @test length(peval) == nrow(deval) + @test all(isfinite, peval) + end @testset "Classification test" begin @@ -114,10 +132,25 @@ end feature_names, ) - # Predictions depend on the number of samples in the dataset + # Test inference output shape and properties + p_cls = m(dtrain) + nclasses = length(levels(dtrain.class)) + @test p_cls isa Matrix{Float32} + @test size(p_cls) == (nrow(dtrain), nclasses) + @test all(isfinite, p_cls) + @test all(sum(p_cls; dims=2) .≈ 1.0) # softmax rows sum to 1 + + # Test eval metric was tracked and decreased + @test haskey(m.info, :logger) + @test !isnothing(m.info[:logger]) + metrics = m.info[:logger][:metrics][:metric] + @test length(metrics) > 1 + @test metrics[end] < metrics[1] # loss decreased from init + + # Test prediction accuracy ptrain = [argmax(x) for x in eachrow(m(dtrain))] peval = [argmax(x) for x in eachrow(m(deval))] @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 @test mean(peval .== levelcode.(deval.class)) > 0.95 -end +end \ No newline at end of file From 2a639542e00045fbcbd0fb0c1dcc97ce3eccc182 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Mon, 16 Feb 2026 19:48:28 +0530 Subject: [PATCH 048/120] revert test Signed-off-by: AdityaPandeyCN --- test/core.jl | 39 +++------------------------------------ 1 file changed, 3 insertions(+), 36 deletions(-) diff --git a/test/core.jl b/test/core.jl index b14e774..6acdc0d 100644 --- a/test/core.jl +++ b/test/core.jl @@ -1,5 +1,5 @@ @testset "Core - data iterators" begin - + end @testset "Core - internals test" begin @@ -36,6 +36,7 @@ end ) m = NeuroTabModel(L, chain, info) + end @testset "Core - Regression" begin @@ -63,7 +64,6 @@ end lr=1e-1, ) - # Test without eval data m = NeuroTabModels.fit( learner, dtrain; @@ -71,13 +71,6 @@ end feature_names ) - # Test inference output - p = m(dtrain) - @test p isa Vector{Float32} - @test length(p) == nrow(dtrain) - @test all(isfinite, p) - - # Test with eval data and early stopping m = NeuroTabModels.fit( learner, dtrain; @@ -86,17 +79,6 @@ end deval, ) - # Test eval metric was tracked - @test haskey(m.info, :logger) - @test !isnothing(m.info[:logger]) - @test length(m.info[:logger][:metrics][:metric]) > 0 - - # Test inference on eval data - peval = m(deval) - @test peval isa Vector{Float32} - @test length(peval) == nrow(deval) - @test all(isfinite, peval) - end @testset "Classification test" begin @@ -132,22 +114,7 @@ end feature_names, ) - # Test inference output shape and properties - p_cls = m(dtrain) - nclasses = length(levels(dtrain.class)) - @test p_cls isa Matrix{Float32} - @test size(p_cls) == (nrow(dtrain), nclasses) - @test all(isfinite, p_cls) - @test all(sum(p_cls; dims=2) .≈ 1.0) # softmax rows sum to 1 - - # Test eval metric was tracked and decreased - @test haskey(m.info, :logger) - @test !isnothing(m.info[:logger]) - metrics = m.info[:logger][:metrics][:metric] - @test length(metrics) > 1 - @test metrics[end] < metrics[1] # loss decreased from init - - # Test prediction accuracy + # Predictions depend on the number of samples in the dataset ptrain = [argmax(x) for x in eachrow(m(dtrain))] peval = [argmax(x) for x in eachrow(m(deval))] @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 From 8f0f70f5a767e17d1928b2a366d79be685daa053 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Tue, 17 Feb 2026 08:45:29 +0530 Subject: [PATCH 049/120] cleanup Signed-off-by: AdityaPandeyCN --- test/core.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/core.jl b/test/core.jl index 6acdc0d..096f250 100644 --- a/test/core.jl +++ b/test/core.jl @@ -120,4 +120,4 @@ end @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 @test mean(peval .== levelcode.(deval.class)) > 0.95 -end \ No newline at end of file +end From d4984a2794a6b27bc2b1df249b41f5b484523875 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Tue, 17 Feb 2026 09:12:19 +0530 Subject: [PATCH 050/120] merge conflicts Signed-off-by: AdityaPandeyCN --- Project.toml | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index b4039a3..9cf7df8 100644 --- a/Project.toml +++ b/Project.toml @@ -12,7 +12,6 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" -Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" @@ -22,11 +21,6 @@ Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -<<<<<<< HEAD -======= -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" -cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" ->>>>>>> upstream/lux-new [compat] CategoricalArrays = "1" @@ -50,4 +44,4 @@ MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test", "MLDatasets", "MLJTestInterface", "MLJBase"] +test = ["Test", "MLDatasets", "MLJTestInterface", "MLJBase"] \ No newline at end of file From c0799f1ed983e1c560ed2b92cb6a6b737bcd9d27 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Tue, 17 Feb 2026 09:14:23 +0530 Subject: [PATCH 051/120] merge conflicts Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 45 +-------------------------------------------- src/MLJ.jl | 8 ++------ 2 files changed, 3 insertions(+), 50 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 6db5e4d..4206c81 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -14,14 +14,8 @@ import Optimisers: OptimiserChain, WeightDecay, NAdam import OneHotArrays: onehotbatch using Lux -<<<<<<< HEAD using Reactant using Lux: cpu_device, reactant_device -======= -using Enzyme, Reactant -# using Zygote -# using Mooncake ->>>>>>> upstream/lux-new using DataFrames using CategoricalArrays @@ -62,13 +56,7 @@ function init( batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) -<<<<<<< HEAD lux_loss = _get_lux_loss(L) -======= - # loss = get_loss_fn(config.loss) - loss = MSELoss() - # loss = BinaryCrossEntropyLoss(; logits=Val(true)), ->>>>>>> upstream/lux-new target_levels = nothing target_isordered = false @@ -82,7 +70,6 @@ function init( outsize = 2 end -<<<<<<< HEAD dtrain_loader = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) all_batches = collect(dtrain_loader) @@ -96,17 +83,6 @@ function init( end data = all_batches |> dev -======= - Reactant.set_default_backend("gpu") - dev = reactant_device() - backend = AutoReactant() - # dev = cpu_device() - # dev = gpu_device() - # backend = AutoZygote() - # backend = AutoMooncake() - - data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev ->>>>>>> upstream/lux-new info = Dict( :nrounds => 0, @@ -123,15 +99,11 @@ function init( opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) -<<<<<<< HEAD cache = Dict( :data => data, :lux_loss => lux_loss, :train_state => ts, ) -======= - cache = (data=data, ts=ts, loss=loss, backend=backend, info=info) ->>>>>>> upstream/lux-new return m, cache end @@ -196,23 +168,8 @@ function fit( (verbosity > 0) && @info "Init training" end -<<<<<<< HEAD while m.info[:nrounds] < config.nrounds fit_iter!(m, cache) -======= - ts = cache[:ts] - while m.info[:nrounds] < config.nrounds - for d in cache[:data] - gs, loss, stats, ts = Training.single_train_step!( - cache.backend, - cache.loss, - (d[1], d[2]), - ts - ) - end - m.info[:nrounds] += 1 - ->>>>>>> upstream/lux-new iter = m.info[:nrounds] if !isnothing(logger) @@ -229,7 +186,7 @@ function fit( _sync_params_to_model!(m, cache) m.info[:logger] = logger - return m, ts + return m end function _sync_params_to_model!(m, cache) diff --git a/src/MLJ.jl b/src/MLJ.jl index 5bf46a4..6a45c4b 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,11 +3,7 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -<<<<<<< HEAD import ..Fit: init, fit_iter!, _sync_params_to_model! -======= -import ..Fit: init ->>>>>>> upstream/lux-new import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema @@ -38,7 +34,7 @@ function fit( fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) while fitresult.info[:nrounds] < model.nrounds - # fit_iter!(fitresult, cache) + fit_iter!(fitresult, cache) end _sync_params_to_model!(fitresult, cache) @@ -64,7 +60,7 @@ function update( ) if okay_to_continue(model, fitresult, cache) while fitresult.info[:nrounds] < model.nrounds - # fit_iter!(fitresult, cache) + fit_iter!(fitresult, cache) end _sync_params_to_model!(fitresult, cache) report = (features=fitresult.info[:feature_names],) From 75c6b53a0f12f5727ccf745c327a17b356623736 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 19 Feb 2026 09:44:05 +0530 Subject: [PATCH 052/120] Switch to lazy data loading to reduce memory usage and replace custom with native for cleaner residual stacking. Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 125 ++++------------ src/NeuroTabModels.jl | 4 +- src/infer.jl | 108 +++++--------- src/losses.jl | 222 +++++++++++++++++------------ src/models/NeuroTree/neurotrees.jl | 49 ++++--- 5 files changed, 227 insertions(+), 281 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 4206c81..642074b 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -11,7 +11,6 @@ using ..Metrics import Random: Xoshiro import MLJModelInterface: fit import Optimisers: OptimiserChain, WeightDecay, NAdam -import OneHotArrays: onehotbatch using Lux using Reactant @@ -29,40 +28,19 @@ function _get_device(config) return reactant_device() end -function _get_lux_loss(L) - if L <: MSE - return MSELoss() - elseif L <: MAE - return MAELoss() - elseif L <: LogLoss - return BinaryCrossEntropyLoss(; logits=Val(true)) - elseif L <: MLogLoss - return CrossEntropyLoss(; logits=Val(true)) - else - return MSELoss() - end -end - -function init( - config::LearnerTypes, - df::AbstractDataFrame; - feature_names, - target_name, - weight_name=nothing, - offset_name=nothing, -) +function init(config::LearnerTypes, df::AbstractDataFrame; feature_names, target_name, weight_name=nothing, offset_name=nothing) dev = _get_device(config) - batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) - lux_loss = _get_lux_loss(L) + lux_loss = get_loss_fn(L) target_levels = nothing target_isordered = false outsize = 1 + if L <: MLogLoss - eltype(df[!, target_name]) <: CategoricalValue || error("Target variable `$target_name` must have its elements `<: CategoricalValue`") + eltype(df[!, target_name]) <: CategoricalValue || error("Target `$target_name` must be `<: CategoricalValue`") target_levels = CategoricalArrays.levels(df[!, target_name]) target_isordered = isordered(df[!, target_name]) outsize = length(target_levels) @@ -70,26 +48,20 @@ function init( outsize = 2 end - dtrain_loader = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) - all_batches = collect(dtrain_loader) + loader = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) - # One-hot encode targets for multiclass classification if L <: MLogLoss - all_batches = map(all_batches) do batch - x = batch[1] - y_oh = Float32.(onehotbatch(vec(batch[2]), UInt32(1):UInt32(outsize))) - length(batch) > 2 ? (x, y_oh, batch[3:end]...) : (x, y_oh) - end + loader = (begin + x, y = b[1], b[2] + y_int = eltype(y) <: Integer ? y : CategoricalArrays.levelcode.(y) + length(b) == 2 ? (x, vec(y_int)) : (x, vec(y_int), b[3:end]...) + end for b in loader) end + + data = (dev(b) for b in loader) - data = all_batches |> dev - - info = Dict( - :nrounds => 0, - :feature_names => feature_names, - :target_levels => target_levels, - :target_isordered => target_isordered, - :device => config.device) + info = Dict(:nrounds => 0, :feature_names => feature_names, :target_levels => target_levels, + :target_isordered => target_isordered, :device => config.device) chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) @@ -99,64 +71,34 @@ function init( opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) - cache = Dict( - :data => data, - :lux_loss => lux_loss, - :train_state => ts, - ) - return m, cache + return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts) end """ - function fit( - config::LearnerTypes, - dtrain; - feature_names, - target_name, - weight_name=nothing, - offset_name=nothing, - deval=nothing, - print_every_n=9999, - verbosity=1, - ) + fit(config::LearnerTypes, dtrain; kwargs...) Training function of NeuroTabModels' internal API. # Arguments - -- `config::LearnerTypes` -- `dtrain`: Must be `<:AbstractDataFrame` - -# Keyword arguments - -- `feature_names`: Required kwarg, a `Vector{Symbol}` or `Vector{String}` of the feature names. -- `target_name`: Required kwarg, a `Symbol` or `String` indicating the name of the target variable. -- `weight_name=nothing` -- `offset_name=nothing` -- `deval=nothing`: Data for tracking evaluation metric and perform early stopping. -- `print_every_n=9999` -- `verbosity=1` +- `config::LearnerTypes`: The configuration object defining the model architecture and training hyperparameters (e.g., `NeuroTabClassifier`, `NeuroTabRegressor`). +- `dtrain`: The training data, must be `<:AbstractDataFrame`. + +# Keyword Arguments +- `feature_names`: Required. A `Vector{Symbol}` or `Vector{String}` of the feature names to use. +- `target_name`: Required. A `Symbol` or `String` indicating the name of the target variable. +- `weight_name=nothing`: Optional `Symbol` or `String` for the sample weights column. +- `offset_name=nothing`: Optional `Symbol` or `String` for the offset column. +- `deval=nothing`: Optional `AbstractDataFrame` for tracking evaluation metrics and performing early stopping. +- `print_every_n=9999`: Integer. Logs training progress every N epochs. +- `verbosity=1`: Integer. Controls the logging level (0 for silent, >0 for info). """ -function fit( - config::LearnerTypes, - dtrain; - feature_names, - target_name, - weight_name=nothing, - offset_name=nothing, - deval=nothing, - print_every_n=9999, - verbosity=1 -) - - feature_names = Symbol.(feature_names) - target_name = Symbol(target_name) +function fit(config::LearnerTypes, dtrain; feature_names, target_name, weight_name=nothing, offset_name=nothing, deval=nothing, print_every_n=9999, verbosity=1) + feature_names, target_name = Symbol.(feature_names), Symbol(target_name) weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) - # initialize callback and logger if tracking eval data logger = nothing if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) @@ -197,15 +139,10 @@ function _sync_params_to_model!(m, cache) end function fit_iter!(m, cache) - ts = cache[:train_state] - lux_loss = cache[:lux_loss] - + ts, lux_loss = cache[:train_state], cache[:lux_loss] for d in cache[:data] - _, loss, _, ts = Training.single_train_step!( - AutoEnzyme(), lux_loss, (d[1], d[2]), ts - ) + _, loss, _, ts = Training.single_train_step!(AutoEnzyme(), lux_loss, d, ts) end - cache[:train_state] = ts m.info[:nrounds] += 1 return nothing diff --git a/src/NeuroTabModels.jl b/src/NeuroTabModels.jl index 6fad70e..ccabb56 100644 --- a/src/NeuroTabModels.jl +++ b/src/NeuroTabModels.jl @@ -9,10 +9,10 @@ include("losses.jl") include("metrics.jl") include("models/models.jl") using .Models -include("infer.jl") -using .Infer include("learners.jl") using .Learners +include("infer.jl") +using .Infer include("Fit/fit.jl") using .Fit include("MLJ.jl") diff --git a/src/infer.jl b/src/infer.jl index 1cd862c..8215126 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -3,11 +3,12 @@ module Infer using ..Data using ..Losses using ..Models +using ..Learners using Lux using Lux: cpu_device, reactant_device using Reactant -using NNlib: sigmoid, softmax! +using NNlib: sigmoid, softmax using DataFrames: AbstractDataFrame import MLUtils: DataLoader @@ -22,96 +23,63 @@ function _get_device(device::Symbol) end end -""" - _setup_infer(m, data, device) -> (ps, st, model_fn, batches) - -""" -function _setup_infer(m::NeuroTabModel, data, device::Symbol) - dev = _get_device(device) - ps = dev(m.info[:ps]) - st = dev(m.info[:st]) - batches = collect(data) |> dev - - if device == :gpu - compiled_model = Reactant.@compile m.chain(first(batches), ps, st) - return ps, st, compiled_model, batches - else - return ps, st, m.chain, batches - end -end - -# Run single forward pass on device, pull result back to CPU -function _forward_batch(model_fn, x, ps, st) - y, _ = model_fn(x, ps, st) - return cpu_device()(y) -end - -function _postprocess(::Type{<:Union{MSE,MAE}}, raw_preds::Vector) +function _postprocess(::Type{<:Union{MSE,MAE}}, raw_preds) return vcat([vec(p) for p in raw_preds]...) end -function _postprocess(::Type{<:LogLoss}, raw_preds::Vector) + +function _postprocess(::Type{<:LogLoss}, raw_preds) p = vcat([vec(p) for p in raw_preds]...) return sigmoid.(p) end -function _postprocess(::Type{<:MLogLoss}, raw_preds::Vector) - # Transpose from [class, batch] -> [batch, class], then softmax - p = vcat([Matrix(p') for p in raw_preds]...) - softmax!(p; dims=2) - return p + +function _postprocess(::Type{<:MLogLoss}, raw_preds) + p_full = reduce(hcat, raw_preds) + p_soft = softmax(p_full; dims=1) + return Matrix(p_soft') end -function _postprocess(::Type{<:GaussianMLE}, raw_preds::Vector) - # Transpose, then exponentiate the σ column (log(σ) -> σ) - p = vcat([Matrix(p') for p in raw_preds]...) - p[:, 2] .= exp.(p[:, 2]) - return p + +function _postprocess(::Type{<:GaussianMLE}, raw_preds) + p_full = reduce(hcat, raw_preds) + p_T = Matrix(p_full') + p_T[:, 2] .= exp.(p_T[:, 2]) + return p_T end -function _postprocess(::Type{<:Tweedie}, raw_preds::Vector) + +function _postprocess(::Type{<:Tweedie}, raw_preds) p = vcat([vec(p) for p in raw_preds]...) return exp.(p) end -""" - infer(m::NeuroTabModel, data::DataLoader; device=:cpu) +function infer(m::NeuroTabModel{L}, data; device=:cpu) where {L} + dev = _get_device(device) + cdev = cpu_device() + ps = dev(m.info[:ps]) + st = dev(m.info[:st]) + + raw_preds = Vector{AbstractArray}() + + for b in data + x = b isa Tuple ? b[1] : b + + x_dev = dev(x) + y_pred, _ = Lux.apply(m.chain, x_dev, ps, st) + push!(raw_preds, cdev(y_pred)) + end -""" -function infer(m::NeuroTabModel{L}, data::DataLoader; device=:cpu) where {L} - ps, st, model_fn, batches = _setup_infer(m, data, device) - raw_preds = [_forward_batch(model_fn, x, ps, st) for x in batches] return _postprocess(L, raw_preds) end -""" - infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu) - -""" function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu) dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) return infer(m, dinfer; device=device) end -""" - (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) - -""" -function (m::NeuroTabModel{L})(x::AbstractMatrix; device=:cpu) where {L} - dev = _get_device(device) - ps, st = dev(m.info[:ps]), dev(m.info[:st]) - x_dev = dev(x) - if device == :gpu - model_fn = Reactant.@compile m.chain(x_dev, ps, st) - else - model_fn = m.chain - end - raw_pred = _forward_batch(model_fn, x_dev, ps, st) - return _postprocess(L, [raw_pred]) -end - -""" - (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) - -""" function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) return infer(m, data; device=device) end -end # module \ No newline at end of file +function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) + return infer(m, [(x,)]; device=device) +end + +end \ No newline at end of file diff --git a/src/losses.jl b/src/losses.jl index 3cbd4bf..7568017 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -3,10 +3,14 @@ module Losses export get_loss_fn, get_loss_type export LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE, Tweedie -import Statistics: mean, std -import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid +import Statistics: mean +import NNlib: logsigmoid, logsoftmax import OneHotArrays: onehotbatch +using Lux +# ----------------------------------------------------------------------- +# Loss Types +# ----------------------------------------------------------------------- abstract type LossType end abstract type MSE <: LossType end abstract type MAE <: LossType end @@ -15,105 +19,123 @@ abstract type MLogLoss <: LossType end abstract type GaussianMLE <: LossType end abstract type Tweedie <: LossType end -function mse(m, x, y) - mean((m(x) .- y) .^ 2) -end -function mse(m, x, y, w) - sum((m(x) .- y) .^ 2 .* w) / sum(w) -end -function mse(m, x, y, w, offset) - sum((m(x) .+ offset .- y) .^ 2 .* w) / sum(w) -end +# ----------------------------------------------------------------------- +# 1. MSE +# ----------------------------------------------------------------------- +struct MSE_Loss <: Lux.AbstractLossFunction end +(::MSE_Loss)(y_pred::AbstractArray, y) = mean((y_pred .- y) .^ 2) +(::MSE_Loss)(y_pred::AbstractArray, y, w) = sum((y_pred .- y) .^ 2 .* w) / sum(w) +(::MSE_Loss)(y_pred::AbstractArray, y, w, offset) = sum((y_pred .+ offset .- y) .^ 2 .* w) / sum(w) -function mae(m, x, y) - mean(abs.(m(x) .- y)) -end -function mae(m, x, y, w) - sum(abs.(m(x) .- y) .* w) / sum(w) -end -function mae(m, x, y, w, offset) - sum(abs.(m(x) .+ offset .- y) .* w) / sum(w) -end +# ----------------------------------------------------------------------- +# 2. MAE +# ----------------------------------------------------------------------- +struct MAE_Loss <: Lux.AbstractLossFunction end +(::MAE_Loss)(y_pred::AbstractArray, y) = mean(abs.(y_pred .- y)) +(::MAE_Loss)(y_pred::AbstractArray, y, w) = sum(abs.(y_pred .- y) .* w) / sum(w) +(::MAE_Loss)(y_pred::AbstractArray, y, w, offset) = sum(abs.(y_pred .+ offset .- y) .* w) / sum(w) -function logloss(m, x, y) - p = m(x) - mean((1 .- y) .* p .- logsigmoid.(p)) -end -function logloss(m, x, y, w) - p = m(x) - sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) -end -function logloss(m, x, y, w, offset) - p = m(x) .+ offset +# ----------------------------------------------------------------------- +# 3. LogLoss +# ----------------------------------------------------------------------- +struct Log_Loss <: Lux.AbstractLossFunction end +function (::Log_Loss)(y_pred::AbstractArray, y) + mean((1 .- y) .* y_pred .- logsigmoid.(y_pred)) +end +function (::Log_Loss)(y_pred::AbstractArray, y, w) + sum(w .* ((1 .- y) .* y_pred .- logsigmoid.(y_pred))) / sum(w) +end +function (::Log_Loss)(y_pred::AbstractArray, y, w, offset) + p = y_pred .+ offset sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end -function tweedie(m, x, y) - rho = eltype(x)(1.5) - p = exp.(m(x)) - mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) -end -function tweedie(m, x, y, w) - rho = eltype(x)(1.5) - p = exp.(m(x)) - sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) / sum(w) -end -function tweedie(m, x, y, w, offset) - rho = eltype(x)(1.5) - p = exp.(m(x) .+ offset) - sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) / sum(w) +# ----------------------------------------------------------------------- +# 4. MLogLoss +# ----------------------------------------------------------------------- +struct MLog_Loss <: Lux.AbstractLossFunction end +function (::MLog_Loss)(y_pred::AbstractArray, y) + k = size(y_pred, 1) + # If y is a Vector, OneHot it. If y is a Matrix, use it as-is. + y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y + p = logsoftmax(y_pred; dims=1) + mean(-sum(y_oh .* p; dims=1)) +end +function (::MLog_Loss)(y_pred::AbstractArray, y, w) + k = size(y_pred, 1) + y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y + p = logsoftmax(y_pred; dims=1) + sum(-sum(y_oh .* p; dims=1) .* w) / sum(w) +end +function (::MLog_Loss)(y_pred::AbstractArray, y, w, offset) + k = size(y_pred, 1) + y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y + p = logsoftmax(y_pred .+ offset; dims=1) + sum(-sum(y_oh .* p; dims=1) .* w) / sum(w) end -function mlogloss(m, x, y) - p = logsoftmax(m(x); dims=1) - k = size(p, 1) - mean(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1)) +# ----------------------------------------------------------------------- +# 5. Tweedie +# ----------------------------------------------------------------------- +struct Tweedie_Loss{T} <: Lux.AbstractLossFunction + rho::T +end +Tweedie_Loss() = Tweedie_Loss(1.5f0) +function (l::Tweedie_Loss)(y_pred::AbstractArray, y) + rho = eltype(y_pred)(l.rho) + p = exp.(y_pred) + term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) + term2 = y .* p .^ (1 - rho) / (1 - rho) + term3 = p .^ (2 - rho) / (2 - rho) + mean(2 .* (term1 .- term2 .+ term3)) +end +function (l::Tweedie_Loss)(y_pred::AbstractArray, y, w) + rho = eltype(y_pred)(l.rho) + p = exp.(y_pred) + term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) + term2 = y .* p .^ (1 - rho) / (1 - rho) + term3 = p .^ (2 - rho) / (2 - rho) + sum(w .* 2 .* (term1 .- term2 .+ term3)) / sum(w) +end +function (l::Tweedie_Loss)(y_pred::AbstractArray, y, w, offset) + rho = eltype(y_pred)(l.rho) + p = exp.(y_pred .+ offset) + term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) + term2 = y .* p .^ (1 - rho) / (1 - rho) + term3 = p .^ (2 - rho) / (2 - rho) + sum(w .* 2 .* (term1 .- term2 .+ term3)) / sum(w) end -function mlogloss(m, x, y, w) - p = logsoftmax(m(x); dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1) .* w) / sum(w) -end -function mlogloss(m, x, y, w, offset) - p = logsoftmax(m(x) .+ offset; dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(vec(y), UInt32(1):UInt32(k)) .* p; dims=1) .* w) / sum(w) -end - -gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}) where {T} = - -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) -gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}, w::AbstractVector{T}) where {T} = - -sum((-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) -function gaussian_mle(m, x, y) - p = m(x) - gaussian_mle_loss(view(p, 1, :), view(p, 2, :), y) -end -function gaussian_mle(m, x, y, w) - p = m(x) - gaussian_mle_loss(view(p, 1, :), view(p, 2, :), y, w) -end -function gaussian_mle(m, x, y, w, offset) - p = m(x) .+ offset - gaussian_mle_loss(view(p, 1, :), view(p, 2, :), y, w) +# ----------------------------------------------------------------------- +# 6. Gaussian MLE +# ----------------------------------------------------------------------- +struct GaussianMLE_Loss <: Lux.AbstractLossFunction end +function (::GaussianMLE_Loss)(y_pred::AbstractArray, y) + μ = view(y_pred, 1, :) + σ = view(y_pred, 2, :) + T = eltype(μ) + loss = -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) + return loss / length(y) +end +function (::GaussianMLE_Loss)(y_pred::AbstractArray, y, w) + μ = view(y_pred, 1, :) + σ = view(y_pred, 2, :) + T = eltype(μ) + elem_loss = -(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) + sum(elem_loss .* w) / sum(w) +end +function (::GaussianMLE_Loss)(y_pred::AbstractArray, y, w, offset) + y_adj = y_pred .+ offset + μ = view(y_adj, 1, :) + σ = view(y_adj, 2, :) + T = eltype(μ) + elem_loss = -(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) + sum(elem_loss .* w) / sum(w) end -const _loss_fn_dict = Dict( - :mse => mse, - :mae => mae, - :logloss => logloss, - :mlogloss => mlogloss, - :gaussian_mle => gaussian_mle, - :tweedie => tweedie, -) - -get_loss_fn(loss::Symbol) = _loss_fn_dict[loss] +# ----------------------------------------------------------------------- +# Mappings +# ----------------------------------------------------------------------- const _loss_type_dict = Dict( :mse => MSE, @@ -126,4 +148,24 @@ const _loss_type_dict = Dict( get_loss_type(loss::Symbol) = _loss_type_dict[loss] +function get_loss_fn(L::Type{<:LossType}) + if L <: MSE + return MSE_Loss() + elseif L <: MAE + return MAE_Loss() + elseif L <: LogLoss + return Log_Loss() + elseif L <: MLogLoss + return MLog_Loss() + elseif L <: GaussianMLE + return GaussianMLE_Loss() + elseif L <: Tweedie + return Tweedie_Loss() + else + return MSE_Loss() + end +end + +get_loss_fn(s::Symbol) = get_loss_fn(get_loss_type(s)) + end \ No newline at end of file diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index f567cd7..32a1832 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -40,14 +40,12 @@ function NeuroTreeConfig(; kwargs...) ) args_ignored = setdiff(keys(kwargs), keys(args)) - args_ignored_str = join(args_ignored, ", ") length(args_ignored) > 0 && - @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(args_ignored_str)." + @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(join(args_ignored, ", "))." args_default = setdiff(keys(args), keys(kwargs)) - args_default_str = join(args_default, ", ") length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(args_default_str)." + @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." args_override = intersect(keys(args), keys(kwargs)) for arg in args_override @@ -71,40 +69,41 @@ function NeuroTreeConfig(; kwargs...) end function (config::NeuroTreeConfig)(; nfeats, outsize) + function build_block(n_in, n_out) + create_tree() = NeuroTree(n_in => n_out; + tree_type=config.tree_type, + depth=config.depth, + trees=config.ntrees, + actA=act_dict[config.actA], + scaler=config.scaler, + init_scale=config.init_scale + ) + + if config.stack_size == 1 + return create_tree() + else + return Parallel(+, [create_tree() for _ in 1:config.stack_size]...) + end + end + if config.MLE_tree_split && outsize == 2 outsize ÷= 2 chain = Chain( BatchNorm(nfeats), Parallel( vcat, - NeuroTree(nfeats => outsize; - tree_type=config.tree_type, - depth=config.depth, - trees=config.ntrees, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale), - NeuroTree(nfeats => outsize; - tree_type=config.tree_type, - depth=config.depth, - trees=config.ntrees, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale), + build_block(nfeats, outsize), + build_block(nfeats, outsize), ) ) else chain = Chain( BatchNorm(nfeats), - NeuroTree(nfeats => outsize; - tree_type=config.tree_type, - depth=config.depth, - trees=config.ntrees, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale), + build_block(nfeats, outsize) ) end + + return chain end function _identity_act(x) From a329289efc863c6c72fbf4b7c2df950f8cc862aa Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 19 Feb 2026 09:57:06 +0530 Subject: [PATCH 053/120] revert Project.toml dependency Signed-off-by: AdityaPandeyCN --- Project.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/Project.toml b/Project.toml index 9cf7df8..2ba2df8 100644 --- a/Project.toml +++ b/Project.toml @@ -9,13 +9,11 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" -MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" -OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" From e704a3839bd28ee14f5639b355a79b724e962ae6 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 19 Feb 2026 14:27:07 +0530 Subject: [PATCH 054/120] revert comments Signed-off-by: AdityaPandeyCN --- src/models/NeuroTree/model.jl | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index f78a54f..0fb93df 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -22,12 +22,13 @@ function NeuroTree((feats, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, Float32(init_scale)) end +# Define the Lux interface function LuxCore.initialparameters(rng::AbstractRNG, l::NeuroTree) return ( - w=Float32.((rand(rng, l.nodes * l.trees, l.feats) .- 0.5) ./ 4), - b=zeros(Float32, l.nodes * l.trees), - s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), - p=Float32.(randn(rng, l.outs, l.leaves * l.trees) .* l.init_scale), + w=Float32.((rand(rng, l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # w + b=zeros(Float32, l.nodes * l.trees), # b + s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), # s + p=Float32.(randn(rng, l.outs, l.leaves * l.trees) .* l.init_scale), # p ) end @@ -40,18 +41,21 @@ end function (l::NeuroTree)(x, ps, st) if l.scaler - nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) + nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] else - nw = (l.actA(ps.w) * x .+ ps.b) + nw = (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] end - nw = reshape(nw, size(st.ml, 2), :) - lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) - lw = reshape(lw, :, size(x, 2)) - y = ps.p * lw ./ l.trees + nw = reshape(nw, size(st.ml, 2), :) # [NT,B] => [N,TB] + lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) # [N,TB] => [L,TB] + lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] + y = ps.p * lw ./ l.trees # [P,LT] * [LT,B] => [P,B] return y, st end +""" + get_logits_mask(::Val{:binary}, depth::Integer) +""" function get_logits_mask(::Val{:binary}, depth::Integer) nodes = 2^depth - 1 leaves = 2^depth @@ -80,6 +84,9 @@ function get_logits_mask(::Val{:oblivious}, depth::Integer) return mask end +""" + get_softplus_mask(::Val{:binary}, depth::Integer) +""" function get_softplus_mask(::Val{:binary}, depth::Integer) nodes = 2^depth - 1 leaves = 2^depth From 20c9b7419809277f970f7c9199a45d970f63f1d4 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Thu, 19 Feb 2026 23:45:28 +0530 Subject: [PATCH 055/120] revert y_pred to p Signed-off-by: AdityaPandeyCN --- src/losses.jl | 102 +++++++++++++++++++++++++------------------------- 1 file changed, 50 insertions(+), 52 deletions(-) diff --git a/src/losses.jl b/src/losses.jl index 7568017..522772a 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -23,30 +23,30 @@ abstract type Tweedie <: LossType end # 1. MSE # ----------------------------------------------------------------------- struct MSE_Loss <: Lux.AbstractLossFunction end -(::MSE_Loss)(y_pred::AbstractArray, y) = mean((y_pred .- y) .^ 2) -(::MSE_Loss)(y_pred::AbstractArray, y, w) = sum((y_pred .- y) .^ 2 .* w) / sum(w) -(::MSE_Loss)(y_pred::AbstractArray, y, w, offset) = sum((y_pred .+ offset .- y) .^ 2 .* w) / sum(w) +(::MSE_Loss)(p::AbstractArray, y) = mean((p .- y) .^ 2) +(::MSE_Loss)(p::AbstractArray, y, w) = sum((p .- y) .^ 2 .* w) / sum(w) +(::MSE_Loss)(p::AbstractArray, y, w, offset) = sum((p .+ offset .- y) .^ 2 .* w) / sum(w) # ----------------------------------------------------------------------- # 2. MAE # ----------------------------------------------------------------------- struct MAE_Loss <: Lux.AbstractLossFunction end -(::MAE_Loss)(y_pred::AbstractArray, y) = mean(abs.(y_pred .- y)) -(::MAE_Loss)(y_pred::AbstractArray, y, w) = sum(abs.(y_pred .- y) .* w) / sum(w) -(::MAE_Loss)(y_pred::AbstractArray, y, w, offset) = sum(abs.(y_pred .+ offset .- y) .* w) / sum(w) +(::MAE_Loss)(p::AbstractArray, y) = mean(abs.(p .- y)) +(::MAE_Loss)(p::AbstractArray, y, w) = sum(abs.(p .- y) .* w) / sum(w) +(::MAE_Loss)(p::AbstractArray, y, w, offset) = sum(abs.(p .+ offset .- y) .* w) / sum(w) # ----------------------------------------------------------------------- # 3. LogLoss # ----------------------------------------------------------------------- struct Log_Loss <: Lux.AbstractLossFunction end -function (::Log_Loss)(y_pred::AbstractArray, y) - mean((1 .- y) .* y_pred .- logsigmoid.(y_pred)) +function (::Log_Loss)(p::AbstractArray, y) + mean((1 .- y) .* p .- logsigmoid.(p)) end -function (::Log_Loss)(y_pred::AbstractArray, y, w) - sum(w .* ((1 .- y) .* y_pred .- logsigmoid.(y_pred))) / sum(w) +function (::Log_Loss)(p::AbstractArray, y, w) + sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end -function (::Log_Loss)(y_pred::AbstractArray, y, w, offset) - p = y_pred .+ offset +function (::Log_Loss)(p::AbstractArray, y, w, offset) + p = p .+ offset sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end @@ -54,55 +54,54 @@ end # 4. MLogLoss # ----------------------------------------------------------------------- struct MLog_Loss <: Lux.AbstractLossFunction end -function (::MLog_Loss)(y_pred::AbstractArray, y) - k = size(y_pred, 1) - # If y is a Vector, OneHot it. If y is a Matrix, use it as-is. +function (::MLog_Loss)(p::AbstractArray, y) + k = size(p, 1) y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y - p = logsoftmax(y_pred; dims=1) - mean(-sum(y_oh .* p; dims=1)) + lp = logsoftmax(p; dims=1) + mean(-sum(y_oh .* lp; dims=1)) end -function (::MLog_Loss)(y_pred::AbstractArray, y, w) - k = size(y_pred, 1) +function (::MLog_Loss)(p::AbstractArray, y, w) + k = size(p, 1) y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y - p = logsoftmax(y_pred; dims=1) - sum(-sum(y_oh .* p; dims=1) .* w) / sum(w) + lp = logsoftmax(p; dims=1) + sum(-sum(y_oh .* lp; dims=1) .* w) / sum(w) end -function (::MLog_Loss)(y_pred::AbstractArray, y, w, offset) - k = size(y_pred, 1) +function (::MLog_Loss)(p::AbstractArray, y, w, offset) + k = size(p, 1) y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y - p = logsoftmax(y_pred .+ offset; dims=1) - sum(-sum(y_oh .* p; dims=1) .* w) / sum(w) + lp = logsoftmax(p .+ offset; dims=1) + sum(-sum(y_oh .* lp; dims=1) .* w) / sum(w) end # ----------------------------------------------------------------------- # 5. Tweedie # ----------------------------------------------------------------------- -struct Tweedie_Loss{T} <: Lux.AbstractLossFunction +struct Tweedie_Loss{T} <: Lux.AbstractLossFunction rho::T end Tweedie_Loss() = Tweedie_Loss(1.5f0) -function (l::Tweedie_Loss)(y_pred::AbstractArray, y) - rho = eltype(y_pred)(l.rho) - p = exp.(y_pred) +function (l::Tweedie_Loss)(p::AbstractArray, y) + rho = eltype(p)(l.rho) + ep = exp.(p) term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) - term2 = y .* p .^ (1 - rho) / (1 - rho) - term3 = p .^ (2 - rho) / (2 - rho) + term2 = y .* ep .^ (1 - rho) / (1 - rho) + term3 = ep .^ (2 - rho) / (2 - rho) mean(2 .* (term1 .- term2 .+ term3)) end -function (l::Tweedie_Loss)(y_pred::AbstractArray, y, w) - rho = eltype(y_pred)(l.rho) - p = exp.(y_pred) +function (l::Tweedie_Loss)(p::AbstractArray, y, w) + rho = eltype(p)(l.rho) + ep = exp.(p) term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) - term2 = y .* p .^ (1 - rho) / (1 - rho) - term3 = p .^ (2 - rho) / (2 - rho) + term2 = y .* ep .^ (1 - rho) / (1 - rho) + term3 = ep .^ (2 - rho) / (2 - rho) sum(w .* 2 .* (term1 .- term2 .+ term3)) / sum(w) end -function (l::Tweedie_Loss)(y_pred::AbstractArray, y, w, offset) - rho = eltype(y_pred)(l.rho) - p = exp.(y_pred .+ offset) +function (l::Tweedie_Loss)(p::AbstractArray, y, w, offset) + rho = eltype(p)(l.rho) + ep = exp.(p .+ offset) term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) - term2 = y .* p .^ (1 - rho) / (1 - rho) - term3 = p .^ (2 - rho) / (2 - rho) + term2 = y .* ep .^ (1 - rho) / (1 - rho) + term3 = ep .^ (2 - rho) / (2 - rho) sum(w .* 2 .* (term1 .- term2 .+ term3)) / sum(w) end @@ -110,24 +109,24 @@ end # 6. Gaussian MLE # ----------------------------------------------------------------------- struct GaussianMLE_Loss <: Lux.AbstractLossFunction end -function (::GaussianMLE_Loss)(y_pred::AbstractArray, y) - μ = view(y_pred, 1, :) - σ = view(y_pred, 2, :) +function (::GaussianMLE_Loss)(p::AbstractArray, y) + μ = view(p, 1, :) + σ = view(p, 2, :) T = eltype(μ) loss = -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) return loss / length(y) end -function (::GaussianMLE_Loss)(y_pred::AbstractArray, y, w) - μ = view(y_pred, 1, :) - σ = view(y_pred, 2, :) +function (::GaussianMLE_Loss)(p::AbstractArray, y, w) + μ = view(p, 1, :) + σ = view(p, 2, :) T = eltype(μ) elem_loss = -(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) sum(elem_loss .* w) / sum(w) end -function (::GaussianMLE_Loss)(y_pred::AbstractArray, y, w, offset) - y_adj = y_pred .+ offset - μ = view(y_adj, 1, :) - σ = view(y_adj, 2, :) +function (::GaussianMLE_Loss)(p::AbstractArray, y, w, offset) + p_adj = p .+ offset + μ = view(p_adj, 1, :) + σ = view(p_adj, 2, :) T = eltype(μ) elem_loss = -(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) sum(elem_loss .* w) / sum(w) @@ -136,7 +135,6 @@ end # ----------------------------------------------------------------------- # Mappings # ----------------------------------------------------------------------- - const _loss_type_dict = Dict( :mse => MSE, :mae => MAE, From 7082adb36b7327cd6350f881849e478aaaff12aa Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 08:36:53 +0530 Subject: [PATCH 056/120] operate callback directly on active object Signed-off-by: AdityaPandeyCN --- src/Fit/callback.jl | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index c063989..410af7f 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -7,6 +7,10 @@ using ..Learners: LearnerTypes using ..Data: get_df_loader_train using ..Metrics +using Lux +using Lux: Training, reactant_device +using Reactant + export CallBack, init_logger, update_logger!, agg_logger struct CallBack{F,D} @@ -14,8 +18,8 @@ struct CallBack{F,D} deval::D end -function (cb::CallBack)(logger, iter, m) - metric = Metrics.get_metric(m, cb.feval, cb.deval) +function (cb::CallBack)(logger, iter, ts::Training.TrainState) + metric = Metrics.get_metric(ts, cb.feval, cb.deval) update_logger!(logger; iter, metric) return nothing end @@ -28,9 +32,14 @@ function CallBack( weight_name=nothing, offset_name=nothing ) + + backend = config.device == :gpu ? "gpu" : "cpu" + Reactant.set_default_backend(backend) + dev = reactant_device() batchsize = config.batchsize feval = metric_dict[config.metric] - deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize) + deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize) |> dev + return CallBack(feval, deval) end @@ -67,11 +76,14 @@ function update_logger!(logger; iter, metric) end function agg_logger(logger_raw::Vector{Dict}) + _l1 = first(logger_raw) best_iters = [d[:best_iter] for d in logger_raw] best_iter = ceil(Int, median(best_iters)) + best_metrics = [d[:best_metric] for d in logger_raw] best_metric = last(best_metrics) + metrics = (layer=Int[], iter=Int[], metric=Float64[]) for i in eachindex(logger_raw) _l = logger_raw[i] @@ -79,16 +91,18 @@ function agg_logger(logger_raw::Vector{Dict}) append!(metrics[:iter], _l[:metrics][:iter]) append!(metrics[:metric], _l[:metrics][:metric]) end + logger = Dict( :name => _l1[:name], :maximise => _l1[:maximise], - :early_stopping_rounds => _l1[:name], + :early_stopping_rounds => _l1[:early_stopping_rounds], :metrics => metrics, :best_iters => best_iters, :best_iter => best_iter, :best_metrics => best_metrics, :best_metric => best_metric, ) + return logger end From dc129e45db4d0846f21bdbdbdf5475f18d03badb Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 09:19:58 +0530 Subject: [PATCH 057/120] cleanup callback.jl Signed-off-by: AdityaPandeyCN --- src/Fit/callback.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 410af7f..f08e651 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -7,9 +7,7 @@ using ..Learners: LearnerTypes using ..Data: get_df_loader_train using ..Metrics -using Lux using Lux: Training, reactant_device -using Reactant export CallBack, init_logger, update_logger!, agg_logger @@ -33,8 +31,6 @@ function CallBack( offset_name=nothing ) - backend = config.device == :gpu ? "gpu" : "cpu" - Reactant.set_default_backend(backend) dev = reactant_device() batchsize = config.batchsize feval = metric_dict[config.metric] From 8abf99bc8c76aa04b18351ee76a37f2920e71419 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 09:43:53 +0530 Subject: [PATCH 058/120] handle MLogLoss with Lux.CrossEntropyLoss Signed-off-by: AdityaPandeyCN --- src/losses.jl | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/src/losses.jl b/src/losses.jl index 522772a..7c3ba49 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -4,8 +4,7 @@ export get_loss_fn, get_loss_type export LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE, Tweedie import Statistics: mean -import NNlib: logsigmoid, logsoftmax -import OneHotArrays: onehotbatch +import NNlib: logsigmoid using Lux # ----------------------------------------------------------------------- @@ -56,21 +55,20 @@ end struct MLog_Loss <: Lux.AbstractLossFunction end function (::MLog_Loss)(p::AbstractArray, y) k = size(p, 1) - y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y - lp = logsoftmax(p; dims=1) - mean(-sum(y_oh .* lp; dims=1)) + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) # (k, batch) + return Lux.CrossEntropyLoss(; logits=Val(true))(p, y_oh) end function (::MLog_Loss)(p::AbstractArray, y, w) k = size(p, 1) - y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y - lp = logsoftmax(p; dims=1) - sum(-sum(y_oh .* lp; dims=1) .* w) / sum(w) + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p, y_oh) + return sum(vec(per_sample) .* vec(w)) / sum(w) end function (::MLog_Loss)(p::AbstractArray, y, w, offset) k = size(p, 1) - y_oh = ndims(y) == 1 ? onehotbatch(vec(y), 1:k) : y - lp = logsoftmax(p .+ offset; dims=1) - sum(-sum(y_oh .* lp; dims=1) .* w) / sum(w) + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p .+ offset, y_oh) + return sum(vec(per_sample) .* vec(w)) / sum(w) end # ----------------------------------------------------------------------- From 6ed2364e458b2e78359c9c3abaacfab1f05410f4 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 09:44:30 +0530 Subject: [PATCH 059/120] use train state for metrics Signed-off-by: AdityaPandeyCN --- src/metrics.jl | 58 +++++++++++++++++++++++++------------------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/src/metrics.jl b/src/metrics.jl index 2efa0e5..8b1b431 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -4,7 +4,8 @@ export metric_dict, is_maximise, get_metric import Statistics: mean, std import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid -import OneHotArrays: onehotbatch +using Lux +using Reactant """ mse(m, x, y; agg=mean) @@ -92,31 +93,25 @@ end mlogloss(m, x, y, w, offset; agg=mean) """ function mlogloss(m, x, y; agg=mean) - p = m(x) - p_t = permutedims(p, (2, 1)) - lsm = logsoftmax(p_t; dims=1) - k = size(lsm, 1) - y_oh = onehotbatch(vec(y), UInt32(1):UInt32(k)) - raw = -sum(y_oh .* lsm; dims=1) - return agg(vec(raw)) + p = m(x) # (k, batch) + k = size(p, 1) + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) # (k, batch) + lsm = logsoftmax(p; dims=1) + return agg(vec(-sum(y_oh .* lsm; dims=1))) end function mlogloss(m, x, y, w; agg=mean) p = m(x) - p_t = permutedims(p, (2, 1)) - lsm = logsoftmax(p_t; dims=1) - k = size(lsm, 1) - y_oh = onehotbatch(vec(y), UInt32(1):UInt32(k)) - raw = -sum(y_oh .* lsm; dims=1) - return agg(vec(raw) .* vec(w)) + k = size(p, 1) + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + lsm = logsoftmax(p; dims=1) + return agg(vec(-sum(y_oh .* lsm; dims=1)) .* vec(w)) end function mlogloss(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - p_t = permutedims(p, (2, 1)) - lsm = logsoftmax(p_t; dims=1) - k = size(lsm, 1) - y_oh = onehotbatch(vec(y), UInt32(1):UInt32(k)) - raw = -sum(y_oh .* lsm; dims=1) - return agg(vec(raw) .* vec(w)) + k = size(p, 1) + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + lsm = logsoftmax(p; dims=1) + return agg(vec(-sum(y_oh .* lsm; dims=1)) .* vec(w)) end gaussian_loss_elt(μ, σ, y) = -σ - (y - μ)^2 / (2 * max(2.0f-7, exp(2 * σ))) @@ -128,24 +123,30 @@ gaussian_loss_elt(μ, σ, y) = -σ - (y - μ)^2 / (2 * max(2.0f-7, exp(2 * σ))) """ function gaussian_mle(m, x, y; agg=mean) p = m(x) - μ = view(p, :, 1) - σ = view(p, :, 2) + μ = view(p, 1, :) + σ = view(p, 2, :) return agg(gaussian_loss_elt.(μ, σ, vec(y))) end function gaussian_mle(m, x, y, w; agg=mean) p = m(x) - μ = view(p, :, 1) - σ = view(p, :, 2) + μ = view(p, 1, :) + σ = view(p, 2, :) return agg(gaussian_loss_elt.(μ, σ, vec(y)) .* vec(w)) end function gaussian_mle(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - μ = view(p, :, 1) - σ = view(p, :, 2) + μ = view(p, 1, :) + σ = view(p, 2, :) return agg(gaussian_loss_elt.(μ, σ, vec(y)) .* vec(w)) end -function get_metric(m, f::Function, data) +function get_metric(ts::Training.TrainState, f::Function, data) + ps = ts.parameters + st_test = Lux.testmode(ts.states) + chain = ts.model + x0 = first(data)[1] + model_compiled = @compile chain(x0, ps, st_test) + m = x -> first(model_compiled(x, ps, st_test)) metric = 0.0f0 ws = 0.0f0 for d in data @@ -156,8 +157,7 @@ function get_metric(m, f::Function, data) ws += size(d[2], ndims(d[2])) end end - metric = metric / ws - return metric + return metric / ws end const metric_dict = Dict( From f72db199ad5ed20cb8a1962e4566c90e346d1760 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 09:51:39 +0530 Subject: [PATCH 060/120] fix data loader Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 48 +++++++++++++++++++++++++++++------------------- 1 file changed, 29 insertions(+), 19 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 642074b..817f8d6 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -28,7 +28,14 @@ function _get_device(config) return reactant_device() end -function init(config::LearnerTypes, df::AbstractDataFrame; feature_names, target_name, weight_name=nothing, offset_name=nothing) +function init( + config::LearnerTypes, + df::AbstractDataFrame; + feature_names, + target_name, + weight_name=nothing, + offset_name=nothing +) dev = _get_device(config) batchsize = config.batchsize nfeats = length(feature_names) @@ -48,20 +55,15 @@ function init(config::LearnerTypes, df::AbstractDataFrame; feature_names, target outsize = 2 end - loader = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) + data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) |> dev - if L <: MLogLoss - loader = (begin - x, y = b[1], b[2] - y_int = eltype(y) <: Integer ? y : CategoricalArrays.levelcode.(y) - length(b) == 2 ? (x, vec(y_int)) : (x, vec(y_int), b[3:end]...) - end for b in loader) - end - - data = (dev(b) for b in loader) - - info = Dict(:nrounds => 0, :feature_names => feature_names, :target_levels => target_levels, - :target_isordered => target_isordered, :device => config.device) + info = Dict( + :nrounds => 0, + :feature_names => feature_names, + :target_levels => target_levels, + :target_isordered => target_isordered, + :device => config.device + ) chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) @@ -92,7 +94,17 @@ Training function of NeuroTabModels' internal API. - `print_every_n=9999`: Integer. Logs training progress every N epochs. - `verbosity=1`: Integer. Controls the logging level (0 for silent, >0 for info). """ -function fit(config::LearnerTypes, dtrain; feature_names, target_name, weight_name=nothing, offset_name=nothing, deval=nothing, print_every_n=9999, verbosity=1) +function fit( + config::LearnerTypes, + dtrain; + feature_names, + target_name, + weight_name=nothing, + offset_name=nothing, + deval=nothing, + print_every_n=9999, + verbosity=1 +) feature_names, target_name = Symbol.(feature_names), Symbol(target_name) weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) @@ -103,8 +115,7 @@ function fit(config::LearnerTypes, dtrain; feature_names, target_name, weight_na if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) - _sync_params_to_model!(m, cache) - cb(logger, 0, m) + cb(logger, 0, cache[:train_state]) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] else (verbosity > 0) && @info "Init training" @@ -115,8 +126,7 @@ function fit(config::LearnerTypes, dtrain; feature_names, target_name, weight_na iter = m.info[:nrounds] if !isnothing(logger) - _sync_params_to_model!(m, cache) - cb(logger, iter, m) + cb(logger, iter, cache[:train_state]) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] end From 9fc003f49c6a162cdd2fb903ddc50e135c922b5f Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 09:56:59 +0530 Subject: [PATCH 061/120] fix fit function arguments Signed-off-by: AdityaPandeyCN --- src/Fit/fit.jl | 38 +++++++++++++++++++++++++++++--------- 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 817f8d6..d18bf96 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -77,22 +77,42 @@ function init( end """ - fit(config::LearnerTypes, dtrain; kwargs...) + function fit( + config::LearnerTypes, + dtrain; + feature_names, + target_name, + weight_name=nothing, + offset_name=nothing, + deval=nothing, + metric=nothing, + print_every_n=9999, + early_stopping_rounds=9999, + verbosity=1, + device=:cpu, + gpuID=0, + ) Training function of NeuroTabModels' internal API. # Arguments -- `config::LearnerTypes`: The configuration object defining the model architecture and training hyperparameters (e.g., `NeuroTabClassifier`, `NeuroTabRegressor`). -- `dtrain`: The training data, must be `<:AbstractDataFrame`. -# Keyword Arguments +- `config::LearnerTypes`: The configuration object defining the model architecture, loss, and training hyperparameters. +- `dtrain`: The training data. Must be `<:AbstractDataFrame`. + +# Keyword arguments + - `feature_names`: Required. A `Vector{Symbol}` or `Vector{String}` of the feature names to use. - `target_name`: Required. A `Symbol` or `String` indicating the name of the target variable. -- `weight_name=nothing`: Optional `Symbol` or `String` for the sample weights column. -- `offset_name=nothing`: Optional `Symbol` or `String` for the offset column. -- `deval=nothing`: Optional `AbstractDataFrame` for tracking evaluation metrics and performing early stopping. -- `print_every_n=9999`: Integer. Logs training progress every N epochs. -- `verbosity=1`: Integer. Controls the logging level (0 for silent, >0 for info). +- `weight_name=nothing`: Optional. A `Symbol` or `String` indicating the sample weights column. +- `offset_name=nothing`: Optional. A `Symbol` or `String` indicating the offset column. +- `deval=nothing`: Optional. Evaluation data (`<:AbstractDataFrame`) for tracking metrics and early stopping. +- `metric=nothing`: Optional. The evaluation metric to track (e.g., `:mse`, `:logloss`). +- `print_every_n=9999`: Integer. Logs training progress to the console every `N` epochs. +- `early_stopping_rounds=9999`: Integer. Stops training if the evaluation metric does not improve for this many rounds. +- `verbosity=1`: Integer. Controls the logging level (`0` for silent, `>0` for info). +- `device=:cpu`: Symbol. Hardware device to use for training (`:cpu` or `:gpu`). +- `gpuID=0`: Integer. Specifies which GPU to use if multiple are available. """ function fit( config::LearnerTypes, From 52a662d5e6fa6802b7d34674fbfe9101d908cc8c Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 11:32:00 +0530 Subject: [PATCH 062/120] fix dimension flow Signed-off-by: AdityaPandeyCN fix dimension flow Signed-off-by: AdityaPandeyCN --- src/models/NeuroTree/neurotrees.jl | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 32a1832..f62b88c 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -70,7 +70,7 @@ end function (config::NeuroTreeConfig)(; nfeats, outsize) function build_block(n_in, n_out) - create_tree() = NeuroTree(n_in => n_out; + create_tree(in_dim, out_dim) = NeuroTree(in_dim => out_dim; tree_type=config.tree_type, depth=config.depth, trees=config.ntrees, @@ -80,10 +80,18 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) ) if config.stack_size == 1 - return create_tree() - else - return Parallel(+, [create_tree() for _ in 1:config.stack_size]...) + return create_tree(n_in, n_out) end + + layers = [create_tree(n_in, config.hidden_size)] + + for _ in 1:(config.stack_size - 2) + push!(layers, SkipConnection(create_tree(config.hidden_size, config.hidden_size), +)) + end + + push!(layers, create_tree(config.hidden_size, n_out)) + + return Chain(layers...) end if config.MLE_tree_split && outsize == 2 @@ -109,10 +117,12 @@ end function _identity_act(x) return x ./ sum(abs.(x), dims=2) end + function _tanh_act(x) x = tanh_fast.(x) return x ./ sum(abs.(x), dims=2) end + function _hardtanh_act(x) x = hardtanh.(x) return x ./ sum(abs.(x), dims=2) From 162094df5a3a792b0f47b44ab0ae872c2d5d871e Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 14:47:45 +0530 Subject: [PATCH 063/120] file changes cleanup Signed-off-by: AdityaPandeyCN --- src/losses.jl | 131 ++++++++++++------------------------------- src/metrics.jl | 147 +++++++++++++++++++++++++------------------------ 2 files changed, 112 insertions(+), 166 deletions(-) diff --git a/src/losses.jl b/src/losses.jl index 7c3ba49..ce58af6 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -7,9 +7,6 @@ import Statistics: mean import NNlib: logsigmoid using Lux -# ----------------------------------------------------------------------- -# Loss Types -# ----------------------------------------------------------------------- abstract type LossType end abstract type MSE <: LossType end abstract type MAE <: LossType end @@ -18,62 +15,24 @@ abstract type MLogLoss <: LossType end abstract type GaussianMLE <: LossType end abstract type Tweedie <: LossType end -# ----------------------------------------------------------------------- -# 1. MSE -# ----------------------------------------------------------------------- struct MSE_Loss <: Lux.AbstractLossFunction end (::MSE_Loss)(p::AbstractArray, y) = mean((p .- y) .^ 2) (::MSE_Loss)(p::AbstractArray, y, w) = sum((p .- y) .^ 2 .* w) / sum(w) (::MSE_Loss)(p::AbstractArray, y, w, offset) = sum((p .+ offset .- y) .^ 2 .* w) / sum(w) -# ----------------------------------------------------------------------- -# 2. MAE -# ----------------------------------------------------------------------- struct MAE_Loss <: Lux.AbstractLossFunction end (::MAE_Loss)(p::AbstractArray, y) = mean(abs.(p .- y)) (::MAE_Loss)(p::AbstractArray, y, w) = sum(abs.(p .- y) .* w) / sum(w) (::MAE_Loss)(p::AbstractArray, y, w, offset) = sum(abs.(p .+ offset .- y) .* w) / sum(w) -# ----------------------------------------------------------------------- -# 3. LogLoss -# ----------------------------------------------------------------------- struct Log_Loss <: Lux.AbstractLossFunction end -function (::Log_Loss)(p::AbstractArray, y) - mean((1 .- y) .* p .- logsigmoid.(p)) -end -function (::Log_Loss)(p::AbstractArray, y, w) - sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) -end +(::Log_Loss)(p::AbstractArray, y) = mean((1 .- y) .* p .- logsigmoid.(p)) +(::Log_Loss)(p::AbstractArray, y, w) = sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) function (::Log_Loss)(p::AbstractArray, y, w, offset) p = p .+ offset sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end -# ----------------------------------------------------------------------- -# 4. MLogLoss -# ----------------------------------------------------------------------- -struct MLog_Loss <: Lux.AbstractLossFunction end -function (::MLog_Loss)(p::AbstractArray, y) - k = size(p, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) # (k, batch) - return Lux.CrossEntropyLoss(; logits=Val(true))(p, y_oh) -end -function (::MLog_Loss)(p::AbstractArray, y, w) - k = size(p, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) - per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p, y_oh) - return sum(vec(per_sample) .* vec(w)) / sum(w) -end -function (::MLog_Loss)(p::AbstractArray, y, w, offset) - k = size(p, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) - per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p .+ offset, y_oh) - return sum(vec(per_sample) .* vec(w)) / sum(w) -end - -# ----------------------------------------------------------------------- -# 5. Tweedie -# ----------------------------------------------------------------------- struct Tweedie_Loss{T} <: Lux.AbstractLossFunction rho::T end @@ -81,86 +40,68 @@ Tweedie_Loss() = Tweedie_Loss(1.5f0) function (l::Tweedie_Loss)(p::AbstractArray, y) rho = eltype(p)(l.rho) ep = exp.(p) - term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) - term2 = y .* ep .^ (1 - rho) / (1 - rho) - term3 = ep .^ (2 - rho) / (2 - rho) - mean(2 .* (term1 .- term2 .+ term3)) + mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) end function (l::Tweedie_Loss)(p::AbstractArray, y, w) rho = eltype(p)(l.rho) ep = exp.(p) - term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) - term2 = y .* ep .^ (1 - rho) / (1 - rho) - term3 = ep .^ (2 - rho) / (2 - rho) - sum(w .* 2 .* (term1 .- term2 .+ term3)) / sum(w) + sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) / sum(w) end function (l::Tweedie_Loss)(p::AbstractArray, y, w, offset) rho = eltype(p)(l.rho) ep = exp.(p .+ offset) - term1 = y .^ (2 - rho) / ((1 - rho) * (2 - rho)) - term2 = y .* ep .^ (1 - rho) / (1 - rho) - term3 = ep .^ (2 - rho) / (2 - rho) - sum(w .* 2 .* (term1 .- term2 .+ term3)) / sum(w) + sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) / sum(w) +end + +struct MLog_Loss <: Lux.AbstractLossFunction end +function (::MLog_Loss)(p::AbstractArray, y) + y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) + Lux.CrossEntropyLoss(; logits=Val(true))(p, y_oh) +end +function (::MLog_Loss)(p::AbstractArray, y, w) + y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) + per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p, y_oh) + sum(vec(per_sample) .* vec(w)) / sum(w) +end +function (::MLog_Loss)(p::AbstractArray, y, w, offset) + y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) + per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p .+ offset, y_oh) + sum(vec(per_sample) .* vec(w)) / sum(w) end -# ----------------------------------------------------------------------- -# 6. Gaussian MLE -# ----------------------------------------------------------------------- struct GaussianMLE_Loss <: Lux.AbstractLossFunction end function (::GaussianMLE_Loss)(p::AbstractArray, y) - μ = view(p, 1, :) - σ = view(p, 2, :) - T = eltype(μ) - loss = -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) - return loss / length(y) + μ, σ, T = view(p, 1, :), view(p, 2, :), eltype(p) + mean(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ))))) end function (::GaussianMLE_Loss)(p::AbstractArray, y, w) - μ = view(p, 1, :) - σ = view(p, 2, :) - T = eltype(μ) - elem_loss = -(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) - sum(elem_loss .* w) / sum(w) + μ, σ, T = view(p, 1, :), view(p, 2, :), eltype(p) + sum(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) end function (::GaussianMLE_Loss)(p::AbstractArray, y, w, offset) p_adj = p .+ offset - μ = view(p_adj, 1, :) - σ = view(p_adj, 2, :) - T = eltype(μ) - elem_loss = -(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) - sum(elem_loss .* w) / sum(w) + μ, σ, T = view(p_adj, 1, :), view(p_adj, 2, :), eltype(p_adj) + sum(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) end -# ----------------------------------------------------------------------- -# Mappings -# ----------------------------------------------------------------------- const _loss_type_dict = Dict( :mse => MSE, :mae => MAE, :logloss => LogLoss, - :tweedie => Tweedie, + :mlogloss => MLogLoss, :gaussian_mle => GaussianMLE, - :mlogloss => MLogLoss + :tweedie => Tweedie, ) get_loss_type(loss::Symbol) = _loss_type_dict[loss] -function get_loss_fn(L::Type{<:LossType}) - if L <: MSE - return MSE_Loss() - elseif L <: MAE - return MAE_Loss() - elseif L <: LogLoss - return Log_Loss() - elseif L <: MLogLoss - return MLog_Loss() - elseif L <: GaussianMLE - return GaussianMLE_Loss() - elseif L <: Tweedie - return Tweedie_Loss() - else - return MSE_Loss() - end -end +get_loss_fn(::Type{<:MSE}) = MSE_Loss() +get_loss_fn(::Type{<:MAE}) = MAE_Loss() +get_loss_fn(::Type{<:LogLoss}) = Log_Loss() +get_loss_fn(::Type{<:MLogLoss}) = MLog_Loss() +get_loss_fn(::Type{<:GaussianMLE}) = GaussianMLE_Loss() +get_loss_fn(::Type{<:Tweedie}) = Tweedie_Loss() +get_loss_fn(::Type{<:LossType}) = MSE_Loss() # fallback get_loss_fn(s::Symbol) = get_loss_fn(get_loss_type(s)) diff --git a/src/metrics.jl b/src/metrics.jl index 8b1b431..a78299f 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -8,145 +8,149 @@ using Lux using Reactant """ - mse(m, x, y; agg=mean) - mse(m, x, y, w; agg=mean) - mse(m, x, y, w, offset; agg=mean) + mse(x, y; agg=mean) + mse(x, y, w; agg=mean) + mse(x, y, w, offset; agg=mean) """ function mse(m, x, y; agg=mean) - return agg((vec(m(x)) .- vec(y)) .^ 2) + metric = agg((vec(m(x)) .- vec(y)) .^ 2) + return metric end function mse(m, x, y, w; agg=mean) - return agg((vec(m(x)) .- vec(y)) .^ 2 .* vec(w)) + metric = agg((vec(m(x)) .- vec(y)) .^ 2 .* vec(w)) + return metric end function mse(m, x, y, w, offset; agg=mean) - return agg((vec(m(x)) .+ vec(offset) .- vec(y)) .^ 2 .* vec(w)) + metric = agg((vec(m(x)) .+ vec(offset) .- vec(y)) .^ 2 .* vec(w)) + return metric end """ - mae(m, x, y; agg=mean) - mae(m, x, y, w; agg=mean) - mae(m, x, y, w, offset; agg=mean) + mae(x, y; agg=mean) + mae(x, y, w; agg=mean) + mae(x, y, w, offset; agg=mean) """ function mae(m, x, y; agg=mean) - return agg(abs.(vec(m(x)) .- vec(y))) + metric = agg(abs.(vec(m(x)) .- vec(y))) + return metric end function mae(m, x, y, w; agg=mean) - return agg(abs.(vec(m(x)) .- vec(y)) .* vec(w)) + metric = agg(abs.(vec(m(x)) .- vec(y)) .* vec(w)) + return metric end function mae(m, x, y, w, offset; agg=mean) - return agg(abs.(vec(m(x)) .+ vec(offset) .- vec(y)) .* vec(w)) + metric = agg(abs.(vec(m(x)) .+ vec(offset) .- vec(y)) .* vec(w)) + return metric end + """ - logloss(m, x, y; agg=mean) - logloss(m, x, y, w; agg=mean) - logloss(m, x, y, w, offset; agg=mean) + logloss(x, y; agg=mean) + logloss(x, y, w; agg=mean) + logloss(x, y, w, offset; agg=mean) """ function logloss(m, x, y; agg=mean) p = vec(m(x)) - y = vec(y) - return agg((1 .- y) .* p .- logsigmoid.(p)) + metric = agg((1 .- vec(y)) .* p .- logsigmoid.(p)) + return metric end function logloss(m, x, y, w; agg=mean) p = vec(m(x)) - y = vec(y) - return agg(((1 .- y) .* p .- logsigmoid.(p)) .* vec(w)) + metric = agg(((1 .- vec(y)) .* p .- logsigmoid.(p)) .* vec(w)) + return metric end function logloss(m, x, y, w, offset; agg=mean) p = vec(m(x)) .+ vec(offset) - y = vec(y) - return agg(((1 .- y) .* p .- logsigmoid.(p)) .* vec(w)) + metric = agg(((1 .- vec(y)) .* p .- logsigmoid.(p)) .* vec(w)) + return metric end + """ - tweedie(m, x, y; agg=mean) - tweedie(m, x, y, w; agg=mean) - tweedie(m, x, y, w, offset; agg=mean) + tweedie(x, y; agg=mean) + tweedie(x, y, w; agg=mean) + tweedie(x, y, w, offset; agg=mean) """ function tweedie(m, x, y; agg=mean) rho = eltype(x)(1.5) p = exp.(vec(m(x))) - y = vec(y) - return agg(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho))) + agg(2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + + p .^ (2 - rho) / (2 - rho)) + ) end -function tweedie(m, x, y, w; agg=mean) +function tweedie(m, x, y, w) + agg = mean rho = eltype(x)(1.5) p = exp.(vec(m(x))) - y = vec(y) - w = vec(w) - return agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho))) + agg(vec(w) .* 2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + + p .^ (2 - rho) / (2 - rho)) + ) end function tweedie(m, x, y, w, offset; agg=mean) rho = eltype(x)(1.5) p = exp.(vec(m(x)) .+ vec(offset)) - y = vec(y) - w = vec(w) - return agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho))) + agg(vec(w) .* 2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + + p .^ (2 - rho) / (2 - rho)) + ) end """ - mlogloss(m, x, y; agg=mean) - mlogloss(m, x, y, w; agg=mean) - mlogloss(m, x, y, w, offset; agg=mean) + mlogloss(x, y; agg=mean) + mlogloss(x, y, w; agg=mean) + mlogloss(x, y, w, offset; agg=mean) """ function mlogloss(m, x, y; agg=mean) - p = m(x) # (k, batch) + p = logsoftmax(m(x); dims=1) k = size(p, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) # (k, batch) - lsm = logsoftmax(p; dims=1) - return agg(vec(-sum(y_oh .* lsm; dims=1))) + raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) + metric = agg(raw) + return metric end function mlogloss(m, x, y, w; agg=mean) - p = m(x) + p = logsoftmax(m(x); dims=1) k = size(p, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) - lsm = logsoftmax(p; dims=1) - return agg(vec(-sum(y_oh .* lsm; dims=1)) .* vec(w)) + raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) + metric = agg(raw .* vec(w)) + return metric end function mlogloss(m, x, y, w, offset; agg=mean) - p = m(x) .+ offset + p = logsoftmax(m(x) .+ offset; dims=1) k = size(p, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) - lsm = logsoftmax(p; dims=1) - return agg(vec(-sum(y_oh .* lsm; dims=1)) .* vec(w)) + raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) + metric = agg(raw .* vec(w)) + return metric end + gaussian_loss_elt(μ, σ, y) = -σ - (y - μ)^2 / (2 * max(2.0f-7, exp(2 * σ))) -""" - gaussian_mle(m, x, y; agg=mean) - gaussian_mle(m, x, y, w; agg=mean) - gaussian_mle(m, x, y, w, offset; agg=mean) + +"""" + gaussian_mle(x, y; agg=mean) + gaussian_mle(x, y, w; agg=mean) + gaussian_mle(x, y, w, offset; agg=mean) """ function gaussian_mle(m, x, y; agg=mean) p = m(x) - μ = view(p, 1, :) - σ = view(p, 2, :) - return agg(gaussian_loss_elt.(μ, σ, vec(y))) + metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y))) + return metric end function gaussian_mle(m, x, y, w; agg=mean) p = m(x) - μ = view(p, 1, :) - σ = view(p, 2, :) - return agg(gaussian_loss_elt.(μ, σ, vec(y)) .* vec(w)) + metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y)) .* vec(w)) + return metric end function gaussian_mle(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - μ = view(p, 1, :) - σ = view(p, 2, :) - return agg(gaussian_loss_elt.(μ, σ, vec(y)) .* vec(w)) + metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y)) .* vec(w)) + return metric end function get_metric(ts::Training.TrainState, f::Function, data) - ps = ts.parameters - st_test = Lux.testmode(ts.states) - chain = ts.model - x0 = first(data)[1] - model_compiled = @compile chain(x0, ps, st_test) - m = x -> first(model_compiled(x, ps, st_test)) + ps, st = ts.parameters, Lux.testmode(ts.states) + model_compiled = @compile ts.model(first(data)[1], ps, st) + m = x -> first(model_compiled(x, ps, st)) + metric = 0.0f0 ws = 0.0f0 for d in data @@ -157,7 +161,8 @@ function get_metric(ts::Training.TrainState, f::Function, data) ws += size(d[2], ndims(d[2])) end end - return metric / ws + metric = metric / ws + return metric end const metric_dict = Dict( From 916de8bf8269d40eb6611a464215cf4d0a67ddd8 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 20 Feb 2026 21:42:08 +0530 Subject: [PATCH 064/120] fix array typing Signed-off-by: AdityaPandeyCN --- src/models/NeuroTree/neurotrees.jl | 31 +++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index f62b88c..b82c28d 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -52,7 +52,7 @@ function NeuroTreeConfig(; kwargs...) args[arg] = kwargs[arg] end - config = NeuroTreeConfig( + return NeuroTreeConfig( Symbol(args[:tree_type]), Symbol(args[:actA]), args[:depth], @@ -64,16 +64,16 @@ function NeuroTreeConfig(; kwargs...) args[:init_scale], args[:MLE_tree_split], ) - - return config end function (config::NeuroTreeConfig)(; nfeats, outsize) function build_block(n_in, n_out) - create_tree(in_dim, out_dim) = NeuroTree(in_dim => out_dim; + create_tree(in_dim, out_dim) = NeuroTree(; + feats=in_dim, + outs=out_dim, tree_type=config.tree_type, depth=config.depth, - trees=config.ntrees, + trees=config.ntrees, actA=act_dict[config.actA], scaler=config.scaler, init_scale=config.init_scale @@ -82,26 +82,27 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) if config.stack_size == 1 return create_tree(n_in, n_out) end - - layers = [create_tree(n_in, config.hidden_size)] - + + layers = Any[create_tree(n_in, config.hidden_size)] + for _ in 1:(config.stack_size - 2) push!(layers, SkipConnection(create_tree(config.hidden_size, config.hidden_size), +)) end - + push!(layers, create_tree(config.hidden_size, n_out)) - + return Chain(layers...) end - if config.MLE_tree_split && outsize == 2 - outsize ÷= 2 + if config.MLE_tree_split + iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") + head_outsize = outsize ÷ 2 chain = Chain( BatchNorm(nfeats), Parallel( vcat, - build_block(nfeats, outsize), - build_block(nfeats, outsize), + build_block(nfeats, head_outsize), + build_block(nfeats, head_outsize), ) ) else @@ -110,7 +111,7 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) build_block(nfeats, outsize) ) end - + return chain end From 62f0a001028e7ba853572dcab18a6e7a8a13e568 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 20 Feb 2026 13:20:09 -0500 Subject: [PATCH 065/120] sync --- benchmarks/benchmark_mse.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 7cac39d..21b9f1c 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -6,7 +6,7 @@ using Random: seed! Threads.nthreads() seed!(123) -nobs = Int(1e6) +nobs = Int(1e5) num_feat = Int(100) @info "testing with: $nobs observations | $num_feat features." X = rand(Float32, nobs, num_feat) From 23854517dfc4cb3c191a51844f54ae02525edb9d Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 20 Feb 2026 13:27:26 -0500 Subject: [PATCH 066/120] up --- benchmarks/YEAR-regression.jl | 10 +++++----- benchmarks/titanic-logloss.jl | 13 +++++-------- 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index e8199ad..ee8c312 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -65,13 +65,13 @@ arch = NeuroTabModels.NeuroTreeConfig(; # MLE_tree_split=false # ) -device = :gpu +device = :cpu loss = :mse # :mse :gaussian_mle :tweedie learner = NeuroTabRegressor( arch; loss, - nrounds=200, + nrounds=20, early_stopping_rounds=2, lr=1e-3, batchsize=1024, @@ -81,18 +81,18 @@ learner = NeuroTabRegressor( m = NeuroTabModels.fit( learner, dtrain; - deval, + # deval, target_name, feature_names, print_every_n=5, ) -p_eval = m(deval; device); +p_eval = m(deval; device=:cpu); p_eval = p_eval[:, 1] mse_eval = mean((p_eval .- deval.y_norm) .^ 2) @info "MSE - deval" mse_eval -p_test = m(dtest; device); +p_test = m(dtest; device=:cpu); p_test = p_test[:, 1] mse_test = mean((p_test .- dtest.y_norm) .^ 2) * std(df_tot.y_raw)^2 @info "MSE - dtest" mse_test diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index c131b9d..8c13303 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -8,9 +8,6 @@ using CategoricalArrays using OrderedCollections using NeuroTabModels -using Reactant -Reactant.set_default_backend("cpu") - Random.seed!(123) df = MLDatasets.Titanic().dataframe @@ -58,14 +55,14 @@ learner = NeuroTabRegressor( @time m = NeuroTabModels.fit( learner, dtrain; - # deval, # removed - inference not yet adapted + deval, # FIXME: important slowdown when deval is used target_name, feature_names, print_every_n=10, ); # commented out - inference not yet adapted -# p_train = m(dtrain) -# p_eval = m(deval) -# @info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) -# @info mean((p_eval .> 0.5) .== (deval[!, target_name] .> 0.5)) +p_train = m(dtrain) +p_eval = m(deval) +@info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) +@info mean((p_eval .> 0.5) .== (deval[!, target_name] .> 0.5)) From 9ab274d70bd9b83d63a956f59f1c0a5ac9929a48 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 20 Feb 2026 13:30:53 -0500 Subject: [PATCH 067/120] up --- benchmarks/benchmark_mse.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 7cac39d..7d6f8d1 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -48,7 +48,7 @@ learner = NeuroTabRegressor( @time m = NeuroTabModels.fit( learner, dtrain; - # deval=dtrain, + # deval=dtrain, # FIXME: very slow when deval is used, need to adapt infer target_name, feature_names, print_every_n=2, From ea2b5de3209c120b03e44a5f0d41516380505472 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sat, 21 Feb 2026 22:54:48 +0530 Subject: [PATCH 068/120] apply partial=false and fix infer Signed-off-by: AdityaPandeyCN --- Project.toml | 5 ++++- src/data.jl | 2 +- src/infer.jl | 21 +++++++++++++++------ 3 files changed, 20 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index 2ba2df8..3d0cae3 100644 --- a/Project.toml +++ b/Project.toml @@ -4,16 +4,19 @@ authors = ["jeremie.db "] version = "0.3.0" [deps] +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -42,4 +45,4 @@ MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test", "MLDatasets", "MLJTestInterface", "MLJBase"] \ No newline at end of file +test = ["Test", "MLDatasets", "MLJTestInterface", "MLJBase"] diff --git a/src/data.jl b/src/data.jl index a6ee8df..4314530 100644 --- a/src/data.jl +++ b/src/data.jl @@ -121,7 +121,7 @@ function get_df_loader_infer( container = ContainerInfer(x, offset) batchsize = min(batchsize, length(container)) - dinfer = DataLoader(container; shuffle=false, batchsize, partial=true, parallel=false) + dinfer = DataLoader(container; shuffle=false, batchsize, partial=false, parallel=false) return dinfer end diff --git a/src/infer.jl b/src/infer.jl index 03df284..dec3104 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -8,6 +8,7 @@ using ..Learners using Lux using Lux: cpu_device, reactant_device using Reactant +using Reactant: @compile using NNlib: sigmoid, softmax using DataFrames: AbstractDataFrame import MLUtils: DataLoader @@ -58,12 +59,20 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu) where {L} raw_preds = Vector{AbstractArray}() - for b in data - x = b isa Tuple ? b[1] : b - - x_dev = dev(x) - y_pred, _ = Lux.apply(m.chain, x_dev, ps, st) - push!(raw_preds, cdev(y_pred)) + if device == :gpu + x0 = let b = first(data); b isa Tuple ? b[1] : b end + model_compiled = @compile m.chain(dev(x0), ps, st) + for b in data + x = b isa Tuple ? b[1] : b + y_pred, _ = model_compiled(dev(x), ps, st) + push!(raw_preds, cdev(y_pred)) + end + else + for b in data + x = b isa Tuple ? b[1] : b + y_pred, _ = Lux.apply(m.chain, x, ps, st) + push!(raw_preds, cdev(y_pred)) + end end return _postprocess(L, raw_preds) From 7461b0f9bc7eb4a61ae1a1688d73b2d0830dc628 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sat, 21 Feb 2026 22:59:22 +0530 Subject: [PATCH 069/120] revert project.toml Signed-off-by: AdityaPandeyCN --- Project.toml | 3 --- benchmarks/benchmark_mse.jl | 4 ++-- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 3d0cae3..23f4c6b 100644 --- a/Project.toml +++ b/Project.toml @@ -4,19 +4,16 @@ authors = ["jeremie.db "] version = "0.3.0" [deps] -BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" -MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" -OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 7d6f8d1..42e145b 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -39,7 +39,7 @@ learner = NeuroTabRegressor( nrounds=10, lr=1e-2, batchsize=2048, - device=:gpu + device=:cpu ) # Reactant GPU: 5.970480 seconds (2.33 M allocations: 5.242 GiB, 3.80% gc time, 0.00% compilation time) @@ -48,7 +48,7 @@ learner = NeuroTabRegressor( @time m = NeuroTabModels.fit( learner, dtrain; - # deval=dtrain, # FIXME: very slow when deval is used, need to adapt infer + deval=dtrain, # FIXME: very slow when deval is used, need to adapt infer target_name, feature_names, print_every_n=2, From 685375db4b512092361e580d8afc1d941404effd Mon Sep 17 00:00:00 2001 From: Jeremie Desgagne-Bouchard Date: Sat, 21 Feb 2026 21:21:17 -0500 Subject: [PATCH 070/120] Jdb/grant review (#24) * Replace Zygote with Enzyme for gradient computation - Add Enzyme.jl dependency - Implement compute_grads() using Enzyme.autodiff with runtime activity mode - Add train/test mode switching to handle BatchNorm mutation issues - Refactor mlogloss to use direct indexing instead of onehotbatch - Configure Enzyme strictAliasing in module __init__ Signed-off-by: AdityaPandeyCN * revert back loss logic Signed-off-by: AdityaPandeyCN * up * up * use reactant and flux approach Signed-off-by: AdityaPandeyCN * formatting Signed-off-by: AdityaPandeyCN * revert changes Signed-off-by: AdityaPandeyCN * add back imports Signed-off-by: AdityaPandeyCN * update callback.jl Signed-off-by: AdityaPandeyCN * revert back NAdam and remove CUDA dependency Signed-off-by: AdityaPandeyCN * Compile gradient function once in init(), reuse across all epochs Signed-off-by: AdityaPandeyCN * remove CuDNN dependency Signed-off-by: AdityaPandeyCN * refactor fit.jl Signed-off-by: AdityaPandeyCN * Add _sync_to_cpu! to MLJ interface to synchronize trained weights from Reactant cache Signed-off-by: AdityaPandeyCN * Update mlogloss to accept pre-encoded one-hot matrices, removing internal onehotbatch calls. Signed-off-by: AdityaPandeyCN * make track_stats false Signed-off-by: AdityaPandeyCN * update benchmark file Signed-off-by: AdityaPandeyCN * simplify cpu sync logic Signed-off-by: AdityaPandeyCN * clean MLJ.jl Signed-off-by: AdityaPandeyCN * fix imports Signed-off-by: AdityaPandeyCN * fix inference, callbacks, and classification Signed-off-by: AdityaPandeyCN Signed-off-by: AdityaPandeyCN * simplify get_device Signed-off-by: AdityaPandeyCN * model cleanup Signed-off-by: AdityaPandeyCN * fix imports Signed-off-by: AdityaPandeyCN * fix imports Signed-off-by: AdityaPandeyCN * add gpu inference support Signed-off-by: AdityaPandeyCN * revert test Signed-off-by: AdityaPandeyCN * cleanup Signed-off-by: AdityaPandeyCN * merge conflicts Signed-off-by: AdityaPandeyCN * merge conflicts Signed-off-by: AdityaPandeyCN * Switch to lazy data loading to reduce memory usage and replace custom with native for cleaner residual stacking. Signed-off-by: AdityaPandeyCN * revert Project.toml dependency Signed-off-by: AdityaPandeyCN * revert comments Signed-off-by: AdityaPandeyCN * revert y_pred to p Signed-off-by: AdityaPandeyCN * operate callback directly on active object Signed-off-by: AdityaPandeyCN * cleanup callback.jl Signed-off-by: AdityaPandeyCN * handle MLogLoss with Lux.CrossEntropyLoss Signed-off-by: AdityaPandeyCN * use train state for metrics Signed-off-by: AdityaPandeyCN * fix data loader Signed-off-by: AdityaPandeyCN * fix fit function arguments Signed-off-by: AdityaPandeyCN * fix dimension flow Signed-off-by: AdityaPandeyCN fix dimension flow Signed-off-by: AdityaPandeyCN * file changes cleanup Signed-off-by: AdityaPandeyCN * fix array typing Signed-off-by: AdityaPandeyCN * sync * up * up * apply partial=false and fix infer Signed-off-by: AdityaPandeyCN * revert project.toml Signed-off-by: AdityaPandeyCN * up * up * bump version --------- Signed-off-by: AdityaPandeyCN Co-authored-by: AdityaPandeyCN --- Project.toml | 18 +--- benchmarks/YEAR-regression.jl | 10 +- benchmarks/benchmark_mse.jl | 6 +- benchmarks/titanic-logloss.jl | 10 +- src/Fit/callback.jl | 15 +-- src/Fit/fit.jl | 118 +++++++++++----------- src/MLJ.jl | 23 ++--- src/NeuroTabModels.jl | 4 +- src/data.jl | 21 +--- src/infer.jl | 142 ++++++++++++-------------- src/losses.jl | 156 +++++++++++++---------------- src/metrics.jl | 76 +++++++------- src/models/NeuroTree/entmax.jl | 1 - src/models/NeuroTree/model.jl | 56 +---------- src/models/NeuroTree/neurotrees.jl | 117 ++++++++-------------- 15 files changed, 321 insertions(+), 452 deletions(-) diff --git a/Project.toml b/Project.toml index 6265a1a..c4fd623 100644 --- a/Project.toml +++ b/Project.toml @@ -1,49 +1,41 @@ name = "NeuroTabModels" uuid = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" authors = ["jeremie.db "] -version = "0.3.0" +version = "0.4.0" [deps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" -Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" -Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" -cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] -CUDA = "5" CategoricalArrays = "1" -ChainRulesCore = "1" DataFrames = "1.3" -Flux = "0.14, 0.15, 0.16" -Functors = "0.5" +Lux = "1" MLJModelInterface = "1.2.1" MLUtils = "0.4" +NNlib = "0.9" Optimisers = "0.4" Random = "1" Statistics = "1" StatsBase = "0.34" Tables = "1.9" -cuDNN = "1" julia = "1.10" [extras] +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d" MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index e8199ad..ee8c312 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -65,13 +65,13 @@ arch = NeuroTabModels.NeuroTreeConfig(; # MLE_tree_split=false # ) -device = :gpu +device = :cpu loss = :mse # :mse :gaussian_mle :tweedie learner = NeuroTabRegressor( arch; loss, - nrounds=200, + nrounds=20, early_stopping_rounds=2, lr=1e-3, batchsize=1024, @@ -81,18 +81,18 @@ learner = NeuroTabRegressor( m = NeuroTabModels.fit( learner, dtrain; - deval, + # deval, target_name, feature_names, print_every_n=5, ) -p_eval = m(deval; device); +p_eval = m(deval; device=:cpu); p_eval = p_eval[:, 1] mse_eval = mean((p_eval .- deval.y_norm) .^ 2) @info "MSE - deval" mse_eval -p_test = m(dtest; device); +p_test = m(dtest; device=:cpu); p_test = p_test[:, 1] mse_test = mean((p_test .- dtest.y_norm) .^ 2) * std(df_tot.y_raw)^2 @info "MSE - dtest" mse_test diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index c1f88d9..eb8379a 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -37,7 +37,6 @@ learner = NeuroTabRegressor( arch; loss=:mse, nrounds=10, - early_stopping_rounds=2, lr=1e-2, batchsize=2048, device=:gpu @@ -46,12 +45,13 @@ learner = NeuroTabRegressor( # Reactant GPU: 5.970480 seconds (2.33 M allocations: 5.242 GiB, 3.80% gc time, 0.00% compilation time) # Zygote GPU: 9.855853 seconds (27.92 M allocations: 6.005 GiB, 3.58% gc time) # 13.557744 seconds (26.40 M allocations: 5.989 GiB, 9.60% gc time) -@time m, ts = NeuroTabModels.fit( +@time m = NeuroTabModels.fit( learner, dtrain; + deval=dtrain, # FIXME: very slow when deval is used, need to adapt infer target_name, feature_names, - print_every_n=10, + print_every_n=2, ); # desktop: 0.771839 seconds (369.20 k allocations: 1.522 GiB, 5.94% gc time) diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 847a292..8c13303 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -55,14 +55,14 @@ learner = NeuroTabRegressor( @time m = NeuroTabModels.fit( learner, dtrain; - # deval, # removed - inference not yet adapted + deval, # FIXME: important slowdown when deval is used target_name, feature_names, print_every_n=10, ); # commented out - inference not yet adapted -# p_train = m(dtrain) -# p_eval = m(deval) -# @info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) -# @info mean((p_eval .> 0.5) .== (deval[!, target_name] .> 0.5)) +p_train = m(dtrain) +p_eval = m(deval) +@info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) +@info mean((p_eval .> 0.5) .== (deval[!, target_name] .> 0.5)) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 276902a..f08e651 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -2,13 +2,13 @@ module CallBacks using DataFrames using Statistics: mean, median -using Flux: cpu, gpu -using CUDA: CuIterator using ..Learners: LearnerTypes using ..Data: get_df_loader_train using ..Metrics +using Lux: Training, reactant_device + export CallBack, init_logger, update_logger!, agg_logger struct CallBack{F,D} @@ -16,8 +16,8 @@ struct CallBack{F,D} deval::D end -function (cb::CallBack)(logger, iter, m) - metric = Metrics.get_metric(m, cb.feval, cb.deval) +function (cb::CallBack)(logger, iter, ts::Training.TrainState) + metric = Metrics.get_metric(ts, cb.feval, cb.deval) update_logger!(logger; iter, metric) return nothing end @@ -31,10 +31,11 @@ function CallBack( offset_name=nothing ) - device = config.device + dev = reactant_device() batchsize = config.batchsize feval = metric_dict[config.metric] - deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize, device) + deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize) |> dev + return CallBack(feval, deval) end @@ -90,7 +91,7 @@ function agg_logger(logger_raw::Vector{Dict}) logger = Dict( :name => _l1[:name], :maximise => _l1[:maximise], - :early_stopping_rounds => _l1[:name], + :early_stopping_rounds => _l1[:early_stopping_rounds], :metrics => metrics, :best_iters => best_iters, :best_iter => best_iter, diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 6522da0..4b4b8ba 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -1,6 +1,6 @@ module Fit -export fit +export fit, fit_iter! using ..Data using ..Learners @@ -10,14 +10,11 @@ using ..Metrics import Random: Xoshiro import MLJModelInterface: fit -import CUDA, cuDNN -import Optimisers -import Optimisers: OptimiserChain, WeightDecay, Adam, NAdam, Nesterov, Descent, Momentum, AdaDelta +import Optimisers: OptimiserChain, WeightDecay, NAdam using Lux -using Enzyme, Reactant -# using Zygote -# using Mooncake +using Reactant +using Lux: cpu_device, reactant_device using DataFrames using CategoricalArrays @@ -25,28 +22,32 @@ using CategoricalArrays include("callback.jl") using .CallBacks +function _get_device(config) + backend = config.device == :gpu ? "gpu" : "cpu" + Reactant.set_default_backend(backend) + return reactant_device() +end + function init( config::LearnerTypes, df::AbstractDataFrame; feature_names, target_name, weight_name=nothing, - offset_name=nothing, + offset_name=nothing ) - - device = config.device + dev = _get_device(config) batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) - # loss = get_loss_fn(config.loss) - loss = MSELoss() - # loss = BinaryCrossEntropyLoss(; logits=Val(true)), + lux_loss = get_loss_fn(L) target_levels = nothing target_isordered = false outsize = 1 + if L <: MLogLoss - eltype(df[!, target_name]) <: CategoricalValue || error("Target variable `$target_name` must have its elements `<: CategoricalValue`") + eltype(df[!, target_name]) <: CategoricalValue || error("Target `$target_name` must be `<: CategoricalValue`") target_levels = CategoricalArrays.levels(df[!, target_name]) target_isordered = isordered(df[!, target_name]) outsize = length(target_levels) @@ -54,38 +55,30 @@ function init( outsize = 2 end - Reactant.set_default_backend("gpu") - dev = reactant_device() - backend = AutoReactant() - # dev = cpu_device() - # dev = gpu_device() - # backend = AutoZygote() - # backend = AutoMooncake() - - data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize, device) |> dev + data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) |> dev info = Dict( :nrounds => 0, :feature_names => feature_names, :target_levels => target_levels, - :target_isordered => target_isordered) + :target_isordered => target_isordered, + :device => config.device + ) chain = config.arch(; nfeats, outsize) m = NeuroTabModel(L, chain, info) - # Parameter and State Variables rng = Xoshiro(config.seed) ps, st = Lux.setup(rng, m.chain) |> dev opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) - cache = (data=data, ts=ts, loss=loss, backend=backend, info=info) - return m, cache + return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts) end """ function fit( - config::NeuroTypes, + config::LearnerTypes, dtrain; feature_names, target_name, @@ -96,27 +89,28 @@ end print_every_n=9999, early_stopping_rounds=9999, verbosity=1, - device=:cpu, - gpuID=0, ) - Training function of NeuroTabModels' internal API. - # Arguments -- `config::LearnerTypes` -- `dtrain`: Must be `<:AbstractDataFrame` +- `config::LearnerTypes`: The configuration object defining the model architecture, loss, and training hyperparameters. +- `dtrain`: The training data. Must be `<:AbstractDataFrame`. # Keyword arguments -- `feature_names`: Required kwarg, a `Vector{Symbol}` or `Vector{String}` of the feature names. -- `target_name` Required kwarg, a `Symbol` or `String` indicating the name of the target variable. -- `weight_name=nothing` -- `offset_name=nothing` -- `deval=nothing` Data for tracking evaluation metric and perform early stopping. -- `print_every_n=9999` -- `verbosity=1` +- `feature_names`: Required. A `Vector{Symbol}` or `Vector{String}` of the feature names to use. +- `target_name`: Required. A `Symbol` or `String` indicating the name of the target variable. +- `weight_name=nothing`: Optional. A `Symbol` or `String` indicating the sample weights column. +- `offset_name=nothing`: Optional. A `Symbol` or `String` indicating the offset column. +- `deval=nothing`: Optional. Evaluation data (`<:AbstractDataFrame`) for tracking metrics and early stopping. +- `metric=nothing`: Optional. The evaluation metric to track (e.g., `:mse`, `:logloss`). +- `print_every_n=9999`: Integer. Logs training progress to the console every `N` epochs. +- `early_stopping_rounds=9999`: Integer. Stops training if the evaluation metric does not improve for this many rounds. +- `verbosity=1`: Integer. Controls the logging level (`0` for silent, `>0` for info). +- `device=:cpu`: Symbol. Hardware device to use for training (`:cpu` or `:gpu`). +- `gpuID=0`: Integer. Specifies which GPU to use if multiple are available. """ + function fit( config::LearnerTypes, dtrain; @@ -128,41 +122,28 @@ function fit( print_every_n=9999, verbosity=1 ) - - feature_names = Symbol.(feature_names) - target_name = Symbol(target_name) + feature_names, target_name = Symbol.(feature_names), Symbol(target_name) weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) - # initialize callback and logger if tracking eval data logger = nothing if !isnothing(deval) cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) - cb(logger, 0, m) + cb(logger, 0, cache[:train_state]) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] else (verbosity > 0) && @info "Init training" end - ts = cache[:ts] while m.info[:nrounds] < config.nrounds - for d in cache[:data] - gs, loss, stats, ts = Training.single_train_step!( - cache.backend, - cache.loss, - (d[1], d[2]), - ts - ) - end - m.info[:nrounds] += 1 - + fit_iter!(m, cache) iter = m.info[:nrounds] if !isnothing(logger) - cb(logger, iter, m) + cb(logger, iter, cache[:train_state]) if verbosity > 0 && iter % print_every_n == 0 @info "iter $iter" metric = logger[:metrics][:metric][end] end @@ -171,8 +152,27 @@ function fit( (verbosity > 0 && iter % print_every_n == 0) && @info "iter $iter" end end + + _sync_params_to_model!(m, cache) m.info[:logger] = logger - return m, ts + return m +end + +function _sync_params_to_model!(m, cache) + ts = cache[:train_state] + cdev = cpu_device() + m.info[:ps] = cdev(ts.parameters) + m.info[:st] = cdev(Lux.testmode(ts.states)) +end + +function fit_iter!(m, cache) + ts, lux_loss = cache[:train_state], cache[:lux_loss] + for d in cache[:data] + _, loss, _, ts = Training.single_train_step!(AutoEnzyme(), lux_loss, d, ts) + end + cache[:train_state] = ts + m.info[:nrounds] += 1 + return nothing end end \ No newline at end of file diff --git a/src/MLJ.jl b/src/MLJ.jl index 9daa549..8ff6b4b 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -3,7 +3,7 @@ module MLJ using Tables using DataFrames import ..Learners: NeuroTabRegressor, NeuroTabClassifier, LearnerTypes -import ..Fit: init +import ..Fit: init, fit_iter!, _sync_params_to_model! import MLJModelInterface as MMI import MLJModelInterface: fit, update, predict, schema @@ -34,9 +34,11 @@ function fit( fitresult, cache = init(model, dtrain; feature_names, target_name, weight_name, offset_name) while fitresult.info[:nrounds] < model.nrounds - # fit_iter!(fitresult, cache) + fit_iter!(fitresult, cache) end + Fit._sync_to_cpu!(fitresult, cache) + _sync_params_to_model!(fitresult, cache) report = (features=fitresult.info[:feature_names],) return fitresult, cache, report end @@ -45,7 +47,6 @@ function okay_to_continue(model, fitresult, cache) return model.nrounds - fitresult.info[:nrounds] >= 0 end -# For EarlyStopping.jl support MMI.iteration_parameter(::Type{<:LearnerTypes}) = :nrounds function update( @@ -59,8 +60,9 @@ function update( ) if okay_to_continue(model, fitresult, cache) while fitresult.info[:nrounds] < model.nrounds - # fit_iter!(fitresult, cache) + fit_iter!(fitresult, cache) end + _sync_params_to_model!(fitresult, cache) report = (features=fitresult.info[:feature_names],) else fitresult, cache, report = fit(model, verbosity, A, y, w) @@ -68,21 +70,18 @@ function update( return fitresult, cache, report end -function predict(::NeuroTabRegressor, fitresult, A; device=:cpu, gpuID=0) - df = DataFrame(A) +function predict(::NeuroTabRegressor, fitresult, A) Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df; device, gpuID) + pred = fitresult(df) return pred end -function predict(::NeuroTabClassifier, fitresult, A; device=:cpu, gpuID=0) - df = DataFrame(A) +function predict(::NeuroTabClassifier, fitresult, A) Tables.istable(A) ? df = DataFrame(A) : error("`A` must be a Table") - pred = fitresult(df; device, gpuID) + pred = fitresult(df) return MMI.UnivariateFinite(fitresult.info[:target_levels], pred) end -# Metadata MMI.metadata_pkg.( (NeuroTabRegressor, NeuroTabClassifier), name="NeuroTabModels", @@ -109,4 +108,4 @@ MMI.metadata_model( path="NeuroTabModels.NeuroTabClassifier", ) -end +end \ No newline at end of file diff --git a/src/NeuroTabModels.jl b/src/NeuroTabModels.jl index 6fad70e..ccabb56 100644 --- a/src/NeuroTabModels.jl +++ b/src/NeuroTabModels.jl @@ -9,10 +9,10 @@ include("losses.jl") include("metrics.jl") include("models/models.jl") using .Models -include("infer.jl") -using .Infer include("learners.jl") using .Learners +include("infer.jl") +using .Infer include("Fit/fit.jl") using .Fit include("MLJ.jl") diff --git a/src/data.jl b/src/data.jl index 7fdd5da..4314530 100644 --- a/src/data.jl +++ b/src/data.jl @@ -7,11 +7,9 @@ import MLUtils: DataLoader using DataFrames using CategoricalArrays -using CUDA: CuIterator """ ContainerTrain - """ struct ContainerTrain{A<:AbstractMatrix,B,C,D} x::A @@ -48,7 +46,6 @@ function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B return (x, y, w, offset) end - function get_df_loader_train( df::AbstractDataFrame; feature_names, @@ -56,8 +53,7 @@ function get_df_loader_train( weight_name=nothing, offset_name=nothing, batchsize, - shuffle=true, - device=:cpu) + shuffle=true) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') @@ -83,10 +79,8 @@ function get_df_loader_train( return dtrain end - """ ContainerInfer - """ struct ContainerInfer{A<:AbstractMatrix,D} x::A @@ -114,8 +108,7 @@ function get_df_loader_infer( df::AbstractDataFrame; feature_names, offset_name=nothing, - batchsize, - device=:cpu) + batchsize) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') @@ -123,17 +116,13 @@ function get_df_loader_infer( offset = if isnothing(offset_name) nothing else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : offset = Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') + isa(offset_name, String) ? Float32.(df[!, offset_name]) : Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') end container = ContainerInfer(x, offset) batchsize = min(batchsize, length(container)) - dinfer = DataLoader(container; shuffle=false, batchsize, partial=true, parallel=false) - if device == :gpu - return CuIterator(dinfer) - else - return dinfer - end + dinfer = DataLoader(container; shuffle=false, batchsize, partial=false, parallel=false) + return dinfer end end #module \ No newline at end of file diff --git a/src/infer.jl b/src/infer.jl index e09ec39..dec3104 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -3,108 +3,92 @@ module Infer using ..Data using ..Losses using ..Models +using ..Learners -using Flux: sigmoid, softmax!, cpu, gpu, onecold +using Lux +using Lux: cpu_device, reactant_device +using Reactant +using Reactant: @compile +using NNlib: sigmoid, softmax using DataFrames: AbstractDataFrame import MLUtils: DataLoader -import CUDA: CuIterator, device! export infer -""" - DL - -Union{NeuroTabModels.CuIterator, NeuroTabModels.DataLoader} -""" -const DL = Union{CuIterator,DataLoader} - -""" -infer(m::NeuroTabModel, data) - -Return the inference of a `NeuroTabModel` over `data`, where `data` is `AbstractDataFrame`. -""" -function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu, gpuID=0) +function _get_device(device::Symbol) if device == :gpu - device!(gpuID) + Reactant.set_default_backend("gpu") + return reactant_device() + else + return cpu_device() end - m = device == :gpu ? m |> gpu : m |> cpu - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048, device) - p = infer(m, dinfer) - return p end +function _postprocess(::Type{<:Union{MSE,MAE}}, raw_preds) + return vcat([vec(p) for p in raw_preds]...) +end -""" - (m::NeuroTabModel)(x::AbstractMatrix) - (m::NeuroTabModel)(data::AbstractDataFrame) - -Inference for NeuroTabModel -""" -function (m::NeuroTabModel)(x::AbstractMatrix) - p = m.chain(x) - if size(p, 1) == 1 - p = dropdims(p; dims=1) - end - return p +function _postprocess(::Type{<:LogLoss}, raw_preds) + p = vcat([vec(p) for p in raw_preds]...) + return sigmoid.(p) end -function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu, gpuID=0) - if device == :gpu - device!(gpuID) - end - m = device == :gpu ? m |> gpu : m |> cpu - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048, device) - p = infer(m, dinfer) - return p + +function _postprocess(::Type{<:MLogLoss}, raw_preds) + p_full = reduce(hcat, raw_preds) + p_soft = softmax(p_full; dims=1) + return Matrix(p_soft') end +function _postprocess(::Type{<:GaussianMLE}, raw_preds) + p_full = reduce(hcat, raw_preds) + p_T = Matrix(p_full') + p_T[:, 2] .= exp.(p_T[:, 2]) + return p_T +end -function infer(m::NeuroTabModel{L}, data::DL) where {L<:Union{MSE,MAE}} - preds = Vector{Float32}[] - for x in data - push!(preds, Vector(m(x))) - end - p = vcat(preds...) - return p +function _postprocess(::Type{<:Tweedie}, raw_preds) + p = vcat([vec(p) for p in raw_preds]...) + return exp.(p) end -function infer(m::NeuroTabModel{<:LogLoss}, data::DL) - preds = Vector{Float32}[] - for x in data - push!(preds, Vector(m(x))) +function infer(m::NeuroTabModel{L}, data; device=:cpu) where {L} + dev = _get_device(device) + cdev = cpu_device() + ps = dev(m.info[:ps]) + st = dev(m.info[:st]) + + raw_preds = Vector{AbstractArray}() + + if device == :gpu + x0 = let b = first(data); b isa Tuple ? b[1] : b end + model_compiled = @compile m.chain(dev(x0), ps, st) + for b in data + x = b isa Tuple ? b[1] : b + y_pred, _ = model_compiled(dev(x), ps, st) + push!(raw_preds, cdev(y_pred)) + end + else + for b in data + x = b isa Tuple ? b[1] : b + y_pred, _ = Lux.apply(m.chain, x, ps, st) + push!(raw_preds, cdev(y_pred)) + end end - p = vcat(preds...) - p .= sigmoid(p) - return p + + return _postprocess(L, raw_preds) end -function infer(m::NeuroTabModel{<:MLogLoss}, data::DL) - preds = Matrix{Float32}[] - for x in data - push!(preds, Matrix(m(x)')) - end - p = vcat(preds...) - softmax!(p; dims=2) - return p +function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu) + dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) + return infer(m, dinfer; device=device) end -function infer(m::NeuroTabModel{<:GaussianMLE}, data::DL) - preds = Matrix{Float32}[] - for x in data - push!(preds, Matrix(m(x)')) - end - p = vcat(preds...) - p[:, 2] .= exp.(p[:, 2]) # reproject log(σ) into σ - return p +function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) + return infer(m, data; device=device) end -function infer(m::NeuroTabModel{L}, data::DL) where {L<:Union{Tweedie}} - preds = Vector{Float32}[] - for x in data - push!(preds, Vector(m(x))) - end - p = vcat(preds...) - p .= exp.(p) - return p +function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) + return infer(m, [(x,)]; device=device) end -end # module +end diff --git a/src/losses.jl b/src/losses.jl index bec273e..ce58af6 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -3,8 +3,9 @@ module Losses export get_loss_fn, get_loss_type export LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE, Tweedie -import Statistics: mean, std -import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid, onehotbatch +import Statistics: mean +import NNlib: logsigmoid +using Lux abstract type LossType end abstract type MSE <: LossType end @@ -14,115 +15,94 @@ abstract type MLogLoss <: LossType end abstract type GaussianMLE <: LossType end abstract type Tweedie <: LossType end -function mse(m, x, y) - mean((m(x) .- y) .^ 2) -end -function mse(m, x, y, w) - sum((m(x) .- y) .^ 2 .* w) / sum(w) -end -function mse(m, x, y, w, offset) - sum((m(x) .+ offset .- y) .^ 2 .* w) / sum(w) -end +struct MSE_Loss <: Lux.AbstractLossFunction end +(::MSE_Loss)(p::AbstractArray, y) = mean((p .- y) .^ 2) +(::MSE_Loss)(p::AbstractArray, y, w) = sum((p .- y) .^ 2 .* w) / sum(w) +(::MSE_Loss)(p::AbstractArray, y, w, offset) = sum((p .+ offset .- y) .^ 2 .* w) / sum(w) -function mae(m, x, y) - mean(abs.(m(x) .- y)) -end -function mae(m, x, y, w) - sum(abs.(m(x) .- y) .* w) / sum(w) -end -function mae(m, x, y, w, offset) - sum(abs.(m(x) .+ offset .- y) .* w) / sum(w) -end +struct MAE_Loss <: Lux.AbstractLossFunction end +(::MAE_Loss)(p::AbstractArray, y) = mean(abs.(p .- y)) +(::MAE_Loss)(p::AbstractArray, y, w) = sum(abs.(p .- y) .* w) / sum(w) +(::MAE_Loss)(p::AbstractArray, y, w, offset) = sum(abs.(p .+ offset .- y) .* w) / sum(w) -function logloss(m, x, y) - p = m(x) - mean((1 .- y) .* p .- logσ.(p)) -end -function logloss(m, x, y, w) - p = m(x) - sum(w .* ((1 .- y) .* p .- logσ.(p))) / sum(w) -end -function logloss(m, x, y, w, offset) - p = m(x) .+ offset - sum(w .* ((1 .- y) .* p .- logσ.(p))) / sum(w) +struct Log_Loss <: Lux.AbstractLossFunction end +(::Log_Loss)(p::AbstractArray, y) = mean((1 .- y) .* p .- logsigmoid.(p)) +(::Log_Loss)(p::AbstractArray, y, w) = sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) +function (::Log_Loss)(p::AbstractArray, y, w, offset) + p = p .+ offset + sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) end -function tweedie(m, x, y) - rho = eltype(x)(1.5) - p = exp.(m(x)) - mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) +struct Tweedie_Loss{T} <: Lux.AbstractLossFunction + rho::T +end +Tweedie_Loss() = Tweedie_Loss(1.5f0) +function (l::Tweedie_Loss)(p::AbstractArray, y) + rho = eltype(p)(l.rho) + ep = exp.(p) + mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) end -function tweedie(m, x, y, w) - rho = eltype(x)(1.5) - p = exp.(m(x)) - sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) / sum(w) +function (l::Tweedie_Loss)(p::AbstractArray, y, w) + rho = eltype(p)(l.rho) + ep = exp.(p) + sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) / sum(w) end -function tweedie(m, x, y, w, offset) - rho = eltype(x)(1.5) - p = exp.(m(x) .+ offset) - sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) / sum(w) +function (l::Tweedie_Loss)(p::AbstractArray, y, w, offset) + rho = eltype(p)(l.rho) + ep = exp.(p .+ offset) + sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) / sum(w) end -function mlogloss(m, x, y) - p = logsoftmax(m(x); dims=1) - k = size(p, 1) - mean(-sum(onehotbatch(y, 1:k) .* p; dims=1)) +struct MLog_Loss <: Lux.AbstractLossFunction end +function (::MLog_Loss)(p::AbstractArray, y) + y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) + Lux.CrossEntropyLoss(; logits=Val(true))(p, y_oh) end -function mlogloss(m, x, y, w) - p = logsoftmax(m(x); dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) +function (::MLog_Loss)(p::AbstractArray, y, w) + y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) + per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p, y_oh) + sum(vec(per_sample) .* vec(w)) / sum(w) end -function mlogloss(m, x, y, w, offset) - p = logsoftmax(m(x) .+ offset; dims=1) - k = size(p, 1) - sum(-sum(onehotbatch(y, 1:k) .* p; dims=1) .* w) / sum(w) +function (::MLog_Loss)(p::AbstractArray, y, w, offset) + y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) + per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p .+ offset, y_oh) + sum(vec(per_sample) .* vec(w)) / sum(w) end -gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}) where {T} = - -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) -gaussian_mle_loss(μ::AbstractVector{T}, σ::AbstractVector{T}, y::AbstractVector{T}, w::AbstractVector{T}) where {T} = - -sum((-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) - -function gaussian_mle(m, x, y) - p = m(x) - gaussian_mle_loss(view(p, 1, :), view(p, 2, :), y) +struct GaussianMLE_Loss <: Lux.AbstractLossFunction end +function (::GaussianMLE_Loss)(p::AbstractArray, y) + μ, σ, T = view(p, 1, :), view(p, 2, :), eltype(p) + mean(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ))))) end -function gaussian_mle(m, x, y, w) - p = m(x) - gaussian_mle_loss(view(p, 1, :), view(p, 2, :), y, w) +function (::GaussianMLE_Loss)(p::AbstractArray, y, w) + μ, σ, T = view(p, 1, :), view(p, 2, :), eltype(p) + sum(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) end -function gaussian_mle(m, x, y, w, offset) - p = m(x) .+ offset - gaussian_mle_loss(view(p, 1, :), view(p, 2, :), y, w) +function (::GaussianMLE_Loss)(p::AbstractArray, y, w, offset) + p_adj = p .+ offset + μ, σ, T = view(p_adj, 1, :), view(p_adj, 2, :), eltype(p_adj) + sum(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) end -const _loss_fn_dict = Dict( - :mse => mse, - :mae => mae, - :logloss => logloss, - :mlogloss => mlogloss, - :gaussian_mle => gaussian_mle, - :tweedie => tweedie, -) - -get_loss_fn(loss::Symbol) = _loss_fn_dict[loss] - const _loss_type_dict = Dict( :mse => MSE, :mae => MAE, :logloss => LogLoss, - :tweedie => Tweedie, + :mlogloss => MLogLoss, :gaussian_mle => GaussianMLE, - :mlogloss => MLogLoss + :tweedie => Tweedie, ) get_loss_type(loss::Symbol) = _loss_type_dict[loss] +get_loss_fn(::Type{<:MSE}) = MSE_Loss() +get_loss_fn(::Type{<:MAE}) = MAE_Loss() +get_loss_fn(::Type{<:LogLoss}) = Log_Loss() +get_loss_fn(::Type{<:MLogLoss}) = MLog_Loss() +get_loss_fn(::Type{<:GaussianMLE}) = GaussianMLE_Loss() +get_loss_fn(::Type{<:Tweedie}) = Tweedie_Loss() +get_loss_fn(::Type{<:LossType}) = MSE_Loss() # fallback + +get_loss_fn(s::Symbol) = get_loss_fn(get_loss_type(s)) + end \ No newline at end of file diff --git a/src/metrics.jl b/src/metrics.jl index 3110010..a78299f 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -1,9 +1,11 @@ module Metrics -export metric_dict, is_maximise +export metric_dict, is_maximise, get_metric import Statistics: mean, std -import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid, onehotbatch +import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid +using Lux +using Reactant """ mse(x, y; agg=mean) @@ -11,15 +13,15 @@ import Flux: logσ, logsoftmax, softmax, relu, hardsigmoid, onehotbatch mse(x, y, w, offset; agg=mean) """ function mse(m, x, y; agg=mean) - metric = agg((m(x) .- y) .^ 2) + metric = agg((vec(m(x)) .- vec(y)) .^ 2) return metric end function mse(m, x, y, w; agg=mean) - metric = agg((m(x) .- y) .^ 2 .* w) + metric = agg((vec(m(x)) .- vec(y)) .^ 2 .* vec(w)) return metric end function mse(m, x, y, w, offset; agg=mean) - metric = agg((m(x) .+ offset .- y) .^ 2 .* w) + metric = agg((vec(m(x)) .+ vec(offset) .- vec(y)) .^ 2 .* vec(w)) return metric end @@ -29,15 +31,15 @@ end mae(x, y, w, offset; agg=mean) """ function mae(m, x, y; agg=mean) - metric = agg(abs.(m(x) .- y)) + metric = agg(abs.(vec(m(x)) .- vec(y))) return metric end function mae(m, x, y, w; agg=mean) - metric = agg(abs.(m(x) .- y) .* w) + metric = agg(abs.(vec(m(x)) .- vec(y)) .* vec(w)) return metric end function mae(m, x, y, w, offset; agg=mean) - metric = agg(abs.(m(x) .+ offset .- y) .* w) + metric = agg(abs.(vec(m(x)) .+ vec(offset) .- vec(y)) .* vec(w)) return metric end @@ -48,18 +50,18 @@ end logloss(x, y, w, offset; agg=mean) """ function logloss(m, x, y; agg=mean) - p = m(x) - metric = agg((1 .- y) .* p .- logσ.(p)) + p = vec(m(x)) + metric = agg((1 .- vec(y)) .* p .- logsigmoid.(p)) return metric end function logloss(m, x, y, w; agg=mean) - p = m(x) - metric = agg(((1 .- y) .* p .- logσ.(p)) .* w) + p = vec(m(x)) + metric = agg(((1 .- vec(y)) .* p .- logsigmoid.(p)) .* vec(w)) return metric end function logloss(m, x, y, w, offset; agg=mean) - p = m(x) .+ offset - metric = agg(((1 .- y) .* p .- logσ.(p)) .* w) + p = vec(m(x)) .+ vec(offset) + metric = agg(((1 .- vec(y)) .* p .- logsigmoid.(p)) .* vec(w)) return metric end @@ -71,23 +73,23 @@ end """ function tweedie(m, x, y; agg=mean) rho = eltype(x)(1.5) - p = exp.(m(x)) - agg(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + + p = exp.(vec(m(x))) + agg(2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + p .^ (2 - rho) / (2 - rho)) ) end function tweedie(m, x, y, w) agg = mean rho = eltype(x)(1.5) - p = exp.(m(x)) - agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + + p = exp.(vec(m(x))) + agg(vec(w) .* 2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + p .^ (2 - rho) / (2 - rho)) ) end function tweedie(m, x, y, w, offset; agg=mean) rho = eltype(x)(1.5) - p = exp.(m(x) .+ offset) - agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* p .^ (1 - rho) / (1 - rho) + + p = exp.(vec(m(x)) .+ vec(offset)) + agg(vec(w) .* 2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + p .^ (2 - rho) / (2 - rho)) ) end @@ -100,34 +102,28 @@ end function mlogloss(m, x, y; agg=mean) p = logsoftmax(m(x); dims=1) k = size(p, 1) - raw = dropdims(-sum(onehotbatch(y, 1:k) .* p; dims=1); dims=1) + raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) metric = agg(raw) return metric end function mlogloss(m, x, y, w; agg=mean) p = logsoftmax(m(x); dims=1) k = size(p, 1) - raw = dropdims(-sum(onehotbatch(y, 1:k) .* p; dims=1); dims=1) - metric = agg(raw .* w) + raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) + metric = agg(raw .* vec(w)) return metric end function mlogloss(m, x, y, w, offset; agg=mean) p = logsoftmax(m(x) .+ offset; dims=1) k = size(p, 1) - raw = dropdims(-sum(onehotbatch(y, 1:k) .* p; dims=1); dims=1) - metric = agg(raw .* w) + raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) + metric = agg(raw .* vec(w)) return metric end -""" - gaussian_mle(μ::T, σ::T, y::T, w::T) where {T<:AbstractFloat} -""" -gaussian_mle(μ::T, σ::T, y::T) where {T<:AbstractFloat} = - -σ - (y - μ)^2 / (2 * max(T(2.0f-7), exp(2 * σ))) +gaussian_loss_elt(μ, σ, y) = -σ - (y - μ)^2 / (2 * max(2.0f-7, exp(2 * σ))) -gaussian_mle(μ::T, σ::T, y::T, w::T) where {T<:AbstractFloat} = - (-σ - (y - μ)^2 / (2 * max(T(2.0f-7), exp(2 * σ)))) * w """" gaussian_mle(x, y; agg=mean) @@ -136,21 +132,25 @@ gaussian_mle(μ::T, σ::T, y::T, w::T) where {T<:AbstractFloat} = """ function gaussian_mle(m, x, y; agg=mean) p = m(x) - metric = agg(gaussian_mle.(view(p, 1, :), view(p, 2, :), y)) + metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y))) return metric end function gaussian_mle(m, x, y, w; agg=mean) p = m(x) - metric = agg(gaussian_mle.(view(p, 1, :), view(p, 2, :), y, w)) + metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y)) .* vec(w)) return metric end function gaussian_mle(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - metric = agg(gaussian_mle.(view(p, 1, :), view(p, 2, :), y, w)) + metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y)) .* vec(w)) return metric end -function get_metric(m, f::Function, data) +function get_metric(ts::Training.TrainState, f::Function, data) + ps, st = ts.parameters, Lux.testmode(ts.states) + model_compiled = @compile ts.model(first(data)[1], ps, st) + m = x -> first(model_compiled(x, ps, st)) + metric = 0.0f0 ws = 0.0f0 for d in data @@ -158,7 +158,7 @@ function get_metric(m, f::Function, data) if length(d) >= 3 ws += sum(d[3]) else - ws += last(size(d[2])) + ws += size(d[2], ndims(d[2])) end end metric = metric / ws @@ -181,4 +181,4 @@ is_maximise(::typeof(mlogloss)) = false is_maximise(::typeof(gaussian_mle)) = true is_maximise(::typeof(tweedie)) = false -end +end \ No newline at end of file diff --git a/src/models/NeuroTree/entmax.jl b/src/models/NeuroTree/entmax.jl index 4f4bc16..77dde49 100644 --- a/src/models/NeuroTree/entmax.jl +++ b/src/models/NeuroTree/entmax.jl @@ -1,6 +1,5 @@ # Reference: https://github.com/deep-spin/entmax/blob/master/entmax/activations.py one_to_vec(x) = reshape(vec(1:size(x, 2)), 1, :) -one_to_vec(x::AnyCuArray) = reshape(CUDA.CuArray(1:size(x, 2)), 1, :) function entmax_threshold_and_support(x) diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 23b3fa2..0fb93df 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -25,17 +25,17 @@ end # Define the Lux interface function LuxCore.initialparameters(rng::AbstractRNG, l::NeuroTree) return ( - w=Float32.((rand(l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # w + w=Float32.((rand(rng, l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # w b=zeros(Float32, l.nodes * l.trees), # b s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), # s - p=Float32.(randn(l.outs, l.leaves * l.trees) .* l.init_scale), # p + p=Float32.(randn(rng, l.outs, l.leaves * l.trees) .* l.init_scale), # p ) end function LuxCore.initialstates(rng::AbstractRNG, l::NeuroTree) return ( - ml=get_logits_mask(Val(l.tree_type), l.depth), - ms=get_softplus_mask(Val(l.tree_type), l.depth) + ml=Float32.(get_logits_mask(Val(l.tree_type), l.depth)), + ms=Float32.(get_softplus_mask(Val(l.tree_type), l.depth)), ) end @@ -105,50 +105,4 @@ function get_softplus_mask(::Val{:oblivious}, depth::Integer) leaves = 2^depth mask = ones(Bool, leaves, depth) return mask -end - - -""" - StackTree -A StackTree is made of a collection of NeuroTree. -""" -struct StackTree - trees::Vector{NeuroTree} -end - -function StackTree((ins, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, depth=4, ntrees=64, proj_size=1, stack_size=1, hidden_size=8, actA=identity, scaler=true, init_scale=1e-1) - @assert stack_size == 1 || hidden_size >= outs - trees = [] - for i in 1:stack_size - if i == 1 - if i < stack_size - tree = NeuroTree(ins => hidden_size; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - else - tree = NeuroTree(ins => outs; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - end - elseif i < stack_size - tree = NeuroTree(hidden_size => hidden_size; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - else - tree = NeuroTree(hidden_size => outs; tree_type, depth, trees, proj_size, actA, scaler, init_scale) - push!(trees, tree) - end - end - m = StackTree(trees) - return m -end - -function (m::StackTree)(x::AbstractMatrix) - p = m.trees[1](x) - for i in 2:length(m.trees) - if i < length(m.trees) - p = p .+ m.trees[i](p) - else - _p = m.trees[i](p) - p = view(p, 1:size(_p, 1), :) .+ _p - end - end - return p -end +end \ No newline at end of file diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 8db3ce1..420096e 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -3,16 +3,9 @@ module NeuroTrees export NeuroTreeConfig using Random -import .Threads: @threads - -# using CUDA -# import Flux -# import Flux: @layer, trainmode!, gradient, Chain, DataLoader, cpu, gpu -# import Flux: relu, logσ, logsoftmax, softmax, softmax!, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch, glorot_uniform -# import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel - using Lux -using NNlib: softplus, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh +using LuxCore +using NNlib: softplus, sigmoid, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture @@ -33,8 +26,6 @@ struct NeuroTreeConfig <: Architecture end function NeuroTreeConfig(; kwargs...) - - # defaults arguments args = Dict{Symbol,Any}( :tree_type => :binary, :actA => :identity, @@ -49,21 +40,19 @@ function NeuroTreeConfig(; kwargs...) ) args_ignored = setdiff(keys(kwargs), keys(args)) - args_ignored_str = join(args_ignored, ", ") length(args_ignored) > 0 && - @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(args_ignored_str)." + @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(join(args_ignored, ", "))." args_default = setdiff(keys(args), keys(kwargs)) - args_default_str = join(args_default, ", ") length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(args_default_str)." + @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." args_override = intersect(keys(args), keys(kwargs)) for arg in args_override args[arg] = kwargs[arg] end - config = NeuroTreeConfig( + return NeuroTreeConfig( Symbol(args[:tree_type]), Symbol(args[:actA]), args[:depth], @@ -75,89 +64,71 @@ function NeuroTreeConfig(; kwargs...) args[:init_scale], args[:MLE_tree_split], ) - - return config end function (config::NeuroTreeConfig)(; nfeats, outsize) + function build_block(n_in, n_out) + create_tree(in_dim, out_dim) = NeuroTree(; + feats=in_dim, + outs=out_dim, + tree_type=config.tree_type, + depth=config.depth, + trees=config.ntrees, + actA=act_dict[config.actA], + scaler=config.scaler, + init_scale=config.init_scale + ) + + if config.stack_size == 1 + return create_tree(n_in, n_out) + end + + layers = Any[create_tree(n_in, config.hidden_size)] + + for _ in 1:(config.stack_size - 2) + push!(layers, SkipConnection(create_tree(config.hidden_size, config.hidden_size), +)) + end + + push!(layers, create_tree(config.hidden_size, n_out)) - if config.MLE_tree_split && outsize == 2 - outsize ÷= 2 + return Chain(layers...) + end + + if config.MLE_tree_split + iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") + head_outsize = outsize ÷ 2 chain = Chain( - BatchNorm(nfeats), + BatchNorm(nfeats, track_stats=false), Parallel( vcat, - StackTree(nfeats => outsize; - tree_type=config.tree_type, - depth=config.depth, - ntrees=config.ntrees, - proj_size=config.proj_size, - stack_size=config.stack_size, - hidden_size=config.hidden_size, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale), - StackTree(nfeats => outsize; - tree_type=config.tree_type, - depth=config.depth, - ntrees=config.ntrees, - proj_size=config.proj_size, - stack_size=config.stack_size, - hidden_size=config.hidden_size, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale) + build_block(nfeats, head_outsize), + build_block(nfeats, head_outsize), ) ) else - # chain = Chain( - # BatchNorm(nfeats), - # StackTree(nfeats => outsize; - # tree_type=config.tree_type, - # depth=config.depth, - # ntrees=config.ntrees, - # proj_size=config.proj_size, - # stack_size=config.stack_size, - # hidden_size=config.hidden_size, - # actA=act_dict[config.actA], - # scaler=config.scaler, - # init_scale=config.init_scale) - # ) chain = Chain( BatchNorm(nfeats), - NeuroTree(nfeats => outsize; - tree_type=config.tree_type, - depth=config.depth, - trees=config.ntrees, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale), + build_block(nfeats, outsize) ) end -end + return chain +end function _identity_act(x) return x ./ sum(abs.(x), dims=2) end + function _tanh_act(x) - x = Flux.tanh_fast.(x) + x = tanh_fast.(x) return x ./ sum(abs.(x), dims=2) end + function _hardtanh_act(x) - x = Flux.hardtanh.(x) + x = hardtanh.(x) return x ./ sum(abs.(x), dims=2) end -""" - act_dict = Dict( - :identity => _identity_act, - :tanh => _tanh_act, - :hardtanh => _hardtanh_act, - ) - -Dictionary mapping features activation name to their function. -""" const act_dict = Dict( :identity => _identity_act, :tanh => _tanh_act, From 309408f2740fa9b4583023760bc84e882614742e Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 21 Feb 2026 21:32:01 -0500 Subject: [PATCH 071/120] up --- benchmarks/benchmark_mse.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index eb8379a..e4f8a51 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -55,5 +55,5 @@ learner = NeuroTabRegressor( ); # desktop: 0.771839 seconds (369.20 k allocations: 1.522 GiB, 5.94% gc time) -# FIXME: need to adapt infer -# @time p_train = m(dtrain; device=:gpu); +# FIXME: need to adapt infer: return only full batches +@time p_train = m(dtrain; device=:gpu); From 491e68d0d46a88f2117c14da53847ab682740ea5 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 21 Feb 2026 21:41:32 -0500 Subject: [PATCH 072/120] cleanup --- benchmarks/benchmark_mse.jl | 9 +++++---- src/Fit/fit.jl | 20 -------------------- 2 files changed, 5 insertions(+), 24 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 5b6d49a..b819565 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -39,7 +39,7 @@ learner = NeuroTabRegressor( nrounds=10, lr=1e-2, batchsize=2048, - device=:cpu + device=:gpu ) # Reactant GPU: 5.970480 seconds (2.33 M allocations: 5.242 GiB, 3.80% gc time, 0.00% compilation time) @@ -48,12 +48,13 @@ learner = NeuroTabRegressor( @time m = NeuroTabModels.fit( learner, dtrain; - deval=dtrain, # FIXME: very slow when deval is used, need to adapt infer + # deval=dtrain, # FIXME: very slow when deval is used / crashed on GPU target_name, feature_names, print_every_n=2, ); -# desktop: 0.771839 seconds (369.20 k allocations: 1.522 GiB, 5.94% gc time) -# FIXME: need to adapt infer: return only full batches +# Reactant CPU: 0.952495 seconds (57.96 k allocations: 1.517 GiB, 0.23% gc time, 0.00% compilation time) +# Reactant CPU: 10.326071 seconds (29.30 k allocations: 13.145 GiB, 1.97% gc time) +# FIXME: need to adapt infer: returns only full batches: length of p_train must be == nrow(dtrain) @time p_train = m(dtrain; device=:gpu); diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 8d38d0a..4b4b8ba 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -174,25 +174,5 @@ function fit_iter!(m, cache) m.info[:nrounds] += 1 return nothing end -# function fit_iter!(m, cache) -# loss, opts, data = cache[:loss], cache[:opts], cache[:dtrain] -# GC.gc(true) -# if typeof(cache[:dtrain]) <: CUDA.CuIterator -# CUDA.reclaim() -# end -# for d in data -# grads = compute_grads(loss, m, d) -# Optimisers.update!(opts, m, grads) -# end -# m.info[:nrounds] += 1 -# return nothing -# end - -# function compute_grads(loss, model, batch) -# dup_model = Duplicated(model) -# const_args = map(Const, batch) -# grads = Flux.gradient((m, args...) -> loss(m, args...), dup_model, const_args...) -# return grads[1] -# end end \ No newline at end of file From 8784a1f338dd7f58c5b6e115d936dbda6fe79d02 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Wed, 25 Feb 2026 21:08:50 -0500 Subject: [PATCH 073/120] up --- benchmarks/titanic-logloss.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 60e73ac..8ef53de 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -41,15 +41,16 @@ arch = NeuroTabModels.NeuroTreeConfig(; stack_size=1, hidden_size=1, actA=:identity, + MLE_tree_split=false, ) learner = NeuroTabRegressor( arch; - loss=:logloss, + loss=:logloss, # FIXME: gaussian_mle don't train nrounds=100, early_stopping_rounds=2, lr=3e-2, - device=:gpu + device=:cpu ) @time m = NeuroTabModels.fit( From 7db9d34e6b37bb4b60c273d7506364a37eb48ba0 Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Sat, 28 Feb 2026 05:01:25 +0530 Subject: [PATCH 074/120] Fix inference and crash from use of deval (#27) * fix inference and deval run Signed-off-by: AdityaPandeyCN * simplify get_metric() Signed-off-by: AdityaPandeyCN * fix get_metric() Signed-off-by: AdityaPandeyCN * set partial=true Signed-off-by: AdityaPandeyCN * fix gaussian_mle function Signed-off-by: AdityaPandeyCN * add custom loss functions Signed-off-by: AdityaPandeyCN * add compile for cpu path Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN --- src/Fit/callback.jl | 45 ++++++++--- src/Fit/fit.jl | 2 +- src/MLJ.jl | 3 +- src/data.jl | 2 +- src/infer.jl | 33 ++++---- src/losses.jl | 179 ++++++++++++++++++++++++++++---------------- src/metrics.jl | 153 ++++++++++++++++++------------------- 7 files changed, 238 insertions(+), 179 deletions(-) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index f08e651..423c4b4 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -7,36 +7,63 @@ using ..Learners: LearnerTypes using ..Data: get_df_loader_train using ..Metrics -using Lux: Training, reactant_device +using Lux: Training, reactant_device, testmode +using Reactant: @compile export CallBack, init_logger, update_logger!, agg_logger -struct CallBack{F,D} - feval::F +struct CallBack{D,C} deval::D + eval_compiled::C end function (cb::CallBack)(logger, iter, ts::Training.TrainState) - metric = Metrics.get_metric(ts, cb.feval, cb.deval) + metric = Metrics.get_metric(ts, cb.deval, cb.eval_compiled) update_logger!(logger; iter, metric) return nothing end function CallBack( config::LearnerTypes, - deval::AbstractDataFrame; + deval::AbstractDataFrame, + ts::Training.TrainState; feature_names, target_name, weight_name=nothing, offset_name=nothing ) - dev = reactant_device() batchsize = config.batchsize feval = metric_dict[config.metric] - deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize) |> dev + deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize, shuffle=false) |> dev + + ps, st = ts.parameters, testmode(ts.states) + d0 = first(deval) + eval_compiled = _compile_eval_step(ts.model, feval, d0, ps, st) - return CallBack(feval, deval) + return CallBack(deval, eval_compiled) +end + +function _compile_eval_step(chain, feval, d0, ps, st) + if length(d0) == 2 + function _step2(x, y, ps, st) + m = x -> first(chain(x, ps, st)) + return feval(m, x, y; agg=sum), eltype(y)(last(size(y))) + end + return @compile _step2(d0[1], d0[2], ps, st) + elseif length(d0) == 3 + function _step3(x, y, w, ps, st) + m = x -> first(chain(x, ps, st)) + return feval(m, x, y, w; agg=sum), sum(w) + end + return @compile _step3(d0[1], d0[2], d0[3], ps, st) + else + function _step4(x, y, w, offset, ps, st) + m = x -> first(chain(x, ps, st)) + return feval(m, x, y, w, offset; agg=sum), sum(w) + end + return @compile _step4(d0[1], d0[2], d0[3], d0[4], ps, st) + end end function init_logger(config::LearnerTypes) @@ -72,7 +99,6 @@ function update_logger!(logger; iter, metric) end function agg_logger(logger_raw::Vector{Dict}) - _l1 = first(logger_raw) best_iters = [d[:best_iter] for d in logger_raw] best_iter = ceil(Int, median(best_iters)) @@ -98,7 +124,6 @@ function agg_logger(logger_raw::Vector{Dict}) :best_metrics => best_metrics, :best_metric => best_metric, ) - return logger end diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 4b4b8ba..cd39c2f 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -130,7 +130,7 @@ function fit( logger = nothing if !isnothing(deval) - cb = CallBack(config, deval; feature_names, target_name, weight_name, offset_name) + cb = CallBack(config, deval, cache[:train_state]; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) cb(logger, 0, cache[:train_state]) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] diff --git a/src/MLJ.jl b/src/MLJ.jl index 8ff6b4b..2288f4b 100644 --- a/src/MLJ.jl +++ b/src/MLJ.jl @@ -36,8 +36,7 @@ function fit( while fitresult.info[:nrounds] < model.nrounds fit_iter!(fitresult, cache) end - Fit._sync_to_cpu!(fitresult, cache) - + _sync_params_to_model!(fitresult, cache) report = (features=fitresult.info[:feature_names],) return fitresult, cache, report diff --git a/src/data.jl b/src/data.jl index 4314530..a6ee8df 100644 --- a/src/data.jl +++ b/src/data.jl @@ -121,7 +121,7 @@ function get_df_loader_infer( container = ContainerInfer(x, offset) batchsize = min(batchsize, length(container)) - dinfer = DataLoader(container; shuffle=false, batchsize, partial=false, parallel=false) + dinfer = DataLoader(container; shuffle=false, batchsize, partial=true, parallel=false) return dinfer end diff --git a/src/infer.jl b/src/infer.jl index dec3104..ff27c94 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -16,12 +16,9 @@ import MLUtils: DataLoader export infer function _get_device(device::Symbol) - if device == :gpu - Reactant.set_default_backend("gpu") - return reactant_device() - else - return cpu_device() - end + backend = device == :gpu ? "gpu" : "cpu" + Reactant.set_default_backend(backend) + return reactant_device() end function _postprocess(::Type{<:Union{MSE,MAE}}, raw_preds) @@ -59,20 +56,18 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu) where {L} raw_preds = Vector{AbstractArray}() - if device == :gpu - x0 = let b = first(data); b isa Tuple ? b[1] : b end - model_compiled = @compile m.chain(dev(x0), ps, st) - for b in data - x = b isa Tuple ? b[1] : b + b_first = first(data) + x0 = b_first isa Tuple ? b_first[1] : b_first + model_compiled = @compile m.chain(dev(x0), ps, st) + + for b in data + x = b isa Tuple ? b[1] : b + if size(x) == size(x0) y_pred, _ = model_compiled(dev(x), ps, st) - push!(raw_preds, cdev(y_pred)) - end - else - for b in data - x = b isa Tuple ? b[1] : b - y_pred, _ = Lux.apply(m.chain, x, ps, st) - push!(raw_preds, cdev(y_pred)) + else + y_pred, _ = Reactant.@jit Lux.apply(m.chain, dev(x), ps, st) end + push!(raw_preds, cdev(y_pred)) end return _postprocess(L, raw_preds) @@ -91,4 +86,4 @@ function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) return infer(m, [(x,)]; device=device) end -end +end \ No newline at end of file diff --git a/src/losses.jl b/src/losses.jl index ce58af6..920b602 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -4,8 +4,7 @@ export get_loss_fn, get_loss_type export LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE, Tweedie import Statistics: mean -import NNlib: logsigmoid -using Lux +import NNlib: logsigmoid, logsoftmax abstract type LossType end abstract type MSE <: LossType end @@ -15,94 +14,142 @@ abstract type MLogLoss <: LossType end abstract type GaussianMLE <: LossType end abstract type Tweedie <: LossType end -struct MSE_Loss <: Lux.AbstractLossFunction end -(::MSE_Loss)(p::AbstractArray, y) = mean((p .- y) .^ 2) -(::MSE_Loss)(p::AbstractArray, y, w) = sum((p .- y) .^ 2 .* w) / sum(w) -(::MSE_Loss)(p::AbstractArray, y, w, offset) = sum((p .+ offset .- y) .^ 2 .* w) / sum(w) - -struct MAE_Loss <: Lux.AbstractLossFunction end -(::MAE_Loss)(p::AbstractArray, y) = mean(abs.(p .- y)) -(::MAE_Loss)(p::AbstractArray, y, w) = sum(abs.(p .- y) .* w) / sum(w) -(::MAE_Loss)(p::AbstractArray, y, w, offset) = sum(abs.(p .+ offset .- y) .* w) / sum(w) +function mse_loss(model, ps, st, data::Tuple{Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred); y = vec(data[2]) + return mean((p .- y) .^ 2), st_, NamedTuple() +end +function mse_loss(model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred); y = vec(data[2]); w = vec(data[3]) + return sum((p .- y) .^ 2 .* w) / sum(w), st_, NamedTuple() +end +function mse_loss(model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred) .+ vec(data[4]); y = vec(data[2]); w = vec(data[3]) + return sum((p .- y) .^ 2 .* w) / sum(w), st_, NamedTuple() +end -struct Log_Loss <: Lux.AbstractLossFunction end -(::Log_Loss)(p::AbstractArray, y) = mean((1 .- y) .* p .- logsigmoid.(p)) -(::Log_Loss)(p::AbstractArray, y, w) = sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) -function (::Log_Loss)(p::AbstractArray, y, w, offset) - p = p .+ offset - sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w) +function mae_loss(model, ps, st, data::Tuple{Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred); y = vec(data[2]) + return mean(abs.(p .- y)), st_, NamedTuple() +end +function mae_loss(model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred); y = vec(data[2]); w = vec(data[3]) + return sum(abs.(p .- y) .* w) / sum(w), st_, NamedTuple() +end +function mae_loss(model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred) .+ vec(data[4]); y = vec(data[2]); w = vec(data[3]) + return sum(abs.(p .- y) .* w) / sum(w), st_, NamedTuple() end -struct Tweedie_Loss{T} <: Lux.AbstractLossFunction - rho::T +function logloss(model, ps, st, data::Tuple{Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred); y = vec(data[2]) + return mean((1 .- y) .* p .- logsigmoid.(p)), st_, NamedTuple() +end +function logloss(model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred); y = vec(data[2]); w = vec(data[3]) + return sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w), st_, NamedTuple() end -Tweedie_Loss() = Tweedie_Loss(1.5f0) -function (l::Tweedie_Loss)(p::AbstractArray, y) - rho = eltype(p)(l.rho) - ep = exp.(p) - mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) +function logloss(model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + p = vec(pred) .+ vec(data[4]); y = vec(data[2]); w = vec(data[3]) + return sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w), st_, NamedTuple() +end + +function mlogloss(model, ps, st, data::Tuple{Any,Any}) + pred, st_ = model(data[1], ps, st) + k = size(pred, 1) + y_oh = (UInt32(1):UInt32(k)) .== reshape(data[2], 1, :) + lsm = logsoftmax(pred; dims=1) + return mean(-sum(y_oh .* lsm; dims=1)), st_, NamedTuple() end -function (l::Tweedie_Loss)(p::AbstractArray, y, w) - rho = eltype(p)(l.rho) - ep = exp.(p) - sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) / sum(w) +function mlogloss(model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + k = size(pred, 1) + y_oh = (UInt32(1):UInt32(k)) .== reshape(data[2], 1, :) + lsm = logsoftmax(pred; dims=1) + return sum(vec(-sum(y_oh .* lsm; dims=1)) .* vec(data[3])) / sum(data[3]), st_, NamedTuple() end -function (l::Tweedie_Loss)(p::AbstractArray, y, w, offset) - rho = eltype(p)(l.rho) - ep = exp.(p .+ offset) - sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) - y .* ep .^ (1 - rho) / (1 - rho) + ep .^ (2 - rho) / (2 - rho))) / sum(w) +function mlogloss(model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + pred = pred .+ data[4] + k = size(pred, 1) + y_oh = (UInt32(1):UInt32(k)) .== reshape(data[2], 1, :) + lsm = logsoftmax(pred; dims=1) + return sum(vec(-sum(y_oh .* lsm; dims=1)) .* vec(data[3])) / sum(data[3]), st_, NamedTuple() end -struct MLog_Loss <: Lux.AbstractLossFunction end -function (::MLog_Loss)(p::AbstractArray, y) - y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) - Lux.CrossEntropyLoss(; logits=Val(true))(p, y_oh) +function tweedie(model, ps, st, data::Tuple{Any,Any}) + pred, st_ = model(data[1], ps, st) + rho = eltype(data[1])(1.5) + ep = exp.(vec(pred)); y = vec(data[2]) + loss = mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ + ep .^ (2 - rho) / (2 - rho))) + return loss, st_, NamedTuple() end -function (::MLog_Loss)(p::AbstractArray, y, w) - y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) - per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p, y_oh) - sum(vec(per_sample) .* vec(w)) / sum(w) +function tweedie(model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + rho = eltype(data[1])(1.5) + ep = exp.(vec(pred)); y = vec(data[2]); w = vec(data[3]) + loss = sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ + ep .^ (2 - rho) / (2 - rho))) / sum(w) + return loss, st_, NamedTuple() end -function (::MLog_Loss)(p::AbstractArray, y, w, offset) - y_oh = (UInt32(1):UInt32(size(p, 1))) .== reshape(y, 1, :) - per_sample = Lux.CrossEntropyLoss(; logits=Val(true), agg=identity)(p .+ offset, y_oh) - sum(vec(per_sample) .* vec(w)) / sum(w) +function tweedie(model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + rho = eltype(data[1])(1.5) + ep = exp.(vec(pred) .+ vec(data[4])); y = vec(data[2]); w = vec(data[3]) + loss = sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ + ep .^ (2 - rho) / (2 - rho))) / sum(w) + return loss, st_, NamedTuple() end -struct GaussianMLE_Loss <: Lux.AbstractLossFunction end -function (::GaussianMLE_Loss)(p::AbstractArray, y) - μ, σ, T = view(p, 1, :), view(p, 2, :), eltype(p) - mean(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ))))) +gaussian_mle_loss(μ::AbstractVector, σ::AbstractVector, y::AbstractVector) = + -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(oftype.(σ, 2e-7), exp.(2 .* σ)))) + +gaussian_mle_loss(μ::AbstractVector, σ::AbstractVector, y::AbstractVector, w::AbstractVector) = + -sum((-σ .- (y .- μ) .^ 2 ./ (2 .* max.(oftype.(σ, 2e-7), exp.(2 .* σ)))) .* w) / sum(w) + +function gaussian_mle(model, ps, st, data::Tuple{Any,Any}) + pred, st_ = model(data[1], ps, st) + loss = gaussian_mle_loss(view(pred, 1, :), view(pred, 2, :), vec(data[2])) + return loss, st_, NamedTuple() end -function (::GaussianMLE_Loss)(p::AbstractArray, y, w) - μ, σ, T = view(p, 1, :), view(p, 2, :), eltype(p) - sum(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) +function gaussian_mle(model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + loss = gaussian_mle_loss(view(pred, 1, :), view(pred, 2, :), vec(data[2]), vec(data[3])) + return loss, st_, NamedTuple() end -function (::GaussianMLE_Loss)(p::AbstractArray, y, w, offset) - p_adj = p .+ offset - μ, σ, T = view(p_adj, 1, :), view(p_adj, 2, :), eltype(p_adj) - sum(-(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(T(2e-7), exp.(2 .* σ)))) .* w) / sum(w) +function gaussian_mle(model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = model(data[1], ps, st) + pred = pred .+ data[4] + loss = gaussian_mle_loss(view(pred, 1, :), view(pred, 2, :), vec(data[2]), vec(data[3])) + return loss, st_, NamedTuple() end +get_loss_fn(::Type{<:MSE}) = mse_loss +get_loss_fn(::Type{<:MAE}) = mae_loss +get_loss_fn(::Type{<:LogLoss}) = logloss +get_loss_fn(::Type{<:MLogLoss}) = mlogloss +get_loss_fn(::Type{<:GaussianMLE}) = gaussian_mle +get_loss_fn(::Type{<:Tweedie}) = tweedie + const _loss_type_dict = Dict( :mse => MSE, :mae => MAE, :logloss => LogLoss, - :mlogloss => MLogLoss, - :gaussian_mle => GaussianMLE, :tweedie => Tweedie, + :gaussian_mle => GaussianMLE, + :mlogloss => MLogLoss, ) get_loss_type(loss::Symbol) = _loss_type_dict[loss] - -get_loss_fn(::Type{<:MSE}) = MSE_Loss() -get_loss_fn(::Type{<:MAE}) = MAE_Loss() -get_loss_fn(::Type{<:LogLoss}) = Log_Loss() -get_loss_fn(::Type{<:MLogLoss}) = MLog_Loss() -get_loss_fn(::Type{<:GaussianMLE}) = GaussianMLE_Loss() -get_loss_fn(::Type{<:Tweedie}) = Tweedie_Loss() -get_loss_fn(::Type{<:LossType}) = MSE_Loss() # fallback - get_loss_fn(s::Symbol) = get_loss_fn(get_loss_type(s)) end \ No newline at end of file diff --git a/src/metrics.jl b/src/metrics.jl index a78299f..ea9e3c4 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -5,164 +5,157 @@ export metric_dict, is_maximise, get_metric import Statistics: mean, std import NNlib: logsigmoid, logsoftmax, softmax, relu, hardsigmoid using Lux -using Reactant """ - mse(x, y; agg=mean) - mse(x, y, w; agg=mean) - mse(x, y, w, offset; agg=mean) + mse(m, x, y; agg=mean) + mse(m, x, y, w; agg=mean) + mse(m, x, y, w, offset; agg=mean) """ function mse(m, x, y; agg=mean) - metric = agg((vec(m(x)) .- vec(y)) .^ 2) - return metric + return agg((vec(m(x)) .- vec(y)) .^ 2) end function mse(m, x, y, w; agg=mean) - metric = agg((vec(m(x)) .- vec(y)) .^ 2 .* vec(w)) - return metric + return agg((vec(m(x)) .- vec(y)) .^ 2 .* vec(w)) end function mse(m, x, y, w, offset; agg=mean) - metric = agg((vec(m(x)) .+ vec(offset) .- vec(y)) .^ 2 .* vec(w)) - return metric + return agg((vec(m(x)) .+ vec(offset) .- vec(y)) .^ 2 .* vec(w)) end """ - mae(x, y; agg=mean) - mae(x, y, w; agg=mean) - mae(x, y, w, offset; agg=mean) + mae(m, x, y; agg=mean) + mae(m, x, y, w; agg=mean) + mae(m, x, y, w, offset; agg=mean) """ function mae(m, x, y; agg=mean) - metric = agg(abs.(vec(m(x)) .- vec(y))) - return metric + return agg(abs.(vec(m(x)) .- vec(y))) end function mae(m, x, y, w; agg=mean) - metric = agg(abs.(vec(m(x)) .- vec(y)) .* vec(w)) - return metric + return agg(abs.(vec(m(x)) .- vec(y)) .* vec(w)) end function mae(m, x, y, w, offset; agg=mean) - metric = agg(abs.(vec(m(x)) .+ vec(offset) .- vec(y)) .* vec(w)) - return metric + return agg(abs.(vec(m(x)) .+ vec(offset) .- vec(y)) .* vec(w)) end - """ - logloss(x, y; agg=mean) - logloss(x, y, w; agg=mean) - logloss(x, y, w, offset; agg=mean) + logloss(m, x, y; agg=mean) + logloss(m, x, y, w; agg=mean) + logloss(m, x, y, w, offset; agg=mean) """ function logloss(m, x, y; agg=mean) p = vec(m(x)) - metric = agg((1 .- vec(y)) .* p .- logsigmoid.(p)) - return metric + y = vec(y) + return agg((1 .- y) .* p .- logsigmoid.(p)) end function logloss(m, x, y, w; agg=mean) p = vec(m(x)) - metric = agg(((1 .- vec(y)) .* p .- logsigmoid.(p)) .* vec(w)) - return metric + y = vec(y) + return agg(((1 .- y) .* p .- logsigmoid.(p)) .* vec(w)) end function logloss(m, x, y, w, offset; agg=mean) p = vec(m(x)) .+ vec(offset) - metric = agg(((1 .- vec(y)) .* p .- logsigmoid.(p)) .* vec(w)) - return metric + y = vec(y) + return agg(((1 .- y) .* p .- logsigmoid.(p)) .* vec(w)) end - """ - tweedie(x, y; agg=mean) - tweedie(x, y, w; agg=mean) - tweedie(x, y, w, offset; agg=mean) + tweedie(m, x, y; agg=mean) + tweedie(m, x, y, w; agg=mean) + tweedie(m, x, y, w, offset; agg=mean) """ function tweedie(m, x, y; agg=mean) rho = eltype(x)(1.5) p = exp.(vec(m(x))) - agg(2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) + y = vec(y) + return agg(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* p .^ (1 - rho) / (1 - rho) .+ + p .^ (2 - rho) / (2 - rho))) end -function tweedie(m, x, y, w) - agg = mean +function tweedie(m, x, y, w; agg=mean) rho = eltype(x)(1.5) p = exp.(vec(m(x))) - agg(vec(w) .* 2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) + y = vec(y) + w = vec(w) + return agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* p .^ (1 - rho) / (1 - rho) .+ + p .^ (2 - rho) / (2 - rho))) end function tweedie(m, x, y, w, offset; agg=mean) rho = eltype(x)(1.5) p = exp.(vec(m(x)) .+ vec(offset)) - agg(vec(w) .* 2 .* (vec(y) .^ (2 - rho) / (1 - rho) / (2 - rho) - vec(y) .* p .^ (1 - rho) / (1 - rho) + - p .^ (2 - rho) / (2 - rho)) - ) + y = vec(y) + w = vec(w) + return agg(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* p .^ (1 - rho) / (1 - rho) .+ + p .^ (2 - rho) / (2 - rho))) end """ - mlogloss(x, y; agg=mean) - mlogloss(x, y, w; agg=mean) - mlogloss(x, y, w, offset; agg=mean) + mlogloss(m, x, y; agg=mean) + mlogloss(m, x, y, w; agg=mean) + mlogloss(m, x, y, w, offset; agg=mean) """ function mlogloss(m, x, y; agg=mean) - p = logsoftmax(m(x); dims=1) + p = m(x) k = size(p, 1) - raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) - metric = agg(raw) - return metric + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + lsm = logsoftmax(p; dims=1) + return agg(vec(-sum(y_oh .* lsm; dims=1))) end function mlogloss(m, x, y, w; agg=mean) - p = logsoftmax(m(x); dims=1) + p = m(x) k = size(p, 1) - raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) - metric = agg(raw .* vec(w)) - return metric + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + lsm = logsoftmax(p; dims=1) + return agg(vec(-sum(y_oh .* lsm; dims=1)) .* vec(w)) end function mlogloss(m, x, y, w, offset; agg=mean) - p = logsoftmax(m(x) .+ offset; dims=1) + p = m(x) .+ offset k = size(p, 1) - raw = vec(-sum(((UInt32(1):UInt32(k)) .== reshape(y, 1, :)) .* p; dims=1)) - metric = agg(raw .* vec(w)) - return metric + y_oh = (UInt32(1):UInt32(k)) .== reshape(y, 1, :) + lsm = logsoftmax(p; dims=1) + return agg(vec(-sum(y_oh .* lsm; dims=1)) .* vec(w)) end +""" + gaussian_mle(m, x, y; agg=mean) + gaussian_mle(m, x, y, w; agg=mean) + gaussian_mle(m, x, y, w, offset; agg=mean) +""" +_gaussian_mle_elt(μ, σ, y) = + -σ - (y - μ)^2 / (2 * max(oftype(σ, 2e-7), exp(2 * σ))) -gaussian_loss_elt(μ, σ, y) = -σ - (y - μ)^2 / (2 * max(2.0f-7, exp(2 * σ))) - +_gaussian_mle_elt(μ, σ, y, w) = + (-σ - (y - μ)^2 / (2 * max(oftype(σ, 2e-7), exp(2 * σ)))) * w -"""" - gaussian_mle(x, y; agg=mean) - gaussian_mle(x, y, w; agg=mean) - gaussian_mle(x, y, w, offset; agg=mean) -""" function gaussian_mle(m, x, y; agg=mean) p = m(x) - metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y))) + metric = agg(_gaussian_mle_elt.(view(p, 1, :), view(p, 2, :), vec(y))) return metric end function gaussian_mle(m, x, y, w; agg=mean) p = m(x) - metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y)) .* vec(w)) + metric = agg(_gaussian_mle_elt.(view(p, 1, :), view(p, 2, :), vec(y), vec(w))) return metric end function gaussian_mle(m, x, y, w, offset; agg=mean) p = m(x) .+ offset - metric = agg(gaussian_loss_elt.(view(p, 1, :), view(p, 2, :), vec(y)) .* vec(w)) + metric = agg(_gaussian_mle_elt.(view(p, 1, :), view(p, 2, :), vec(y), vec(w))) return metric end -function get_metric(ts::Training.TrainState, f::Function, data) - ps, st = ts.parameters, Lux.testmode(ts.states) - model_compiled = @compile ts.model(first(data)[1], ps, st) - m = x -> first(model_compiled(x, ps, st)) - +function get_metric(ts, data, eval_compiled) metric = 0.0f0 ws = 0.0f0 + st = Lux.testmode(ts.states) for d in data - metric += f(m, d...; agg=sum) - if length(d) >= 3 - ws += sum(d[3]) + if length(d) == 2 + m_val, w_val = eval_compiled(d[1], d[2], ts.parameters, st) + elseif length(d) == 3 + m_val, w_val = eval_compiled(d[1], d[2], d[3], ts.parameters, st) else - ws += size(d[2], ndims(d[2])) + m_val, w_val = eval_compiled(d[1], d[2], d[3], d[4], ts.parameters, st) end + metric += Float32(m_val) + ws += Float32(w_val) end - metric = metric / ws - return metric + return Float64(metric / ws) end const metric_dict = Dict( From 69641bcdcfbc9a02bd3dfba8cc0378aec8ddab31 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Fri, 27 Feb 2026 22:32:19 -0500 Subject: [PATCH 075/120] up --- src/models/NeuroTree/neurotrees.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 68e450e..6ccef12 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -119,11 +119,11 @@ function _identity_act(x) return x ./ sum(abs.(x), dims=2) end function _tanh_act(x) - x = Flux.tanh_fast.(x) + x = tanh_fast.(x) return x ./ sum(abs.(x), dims=2) end function _hardtanh_act(x) - x = Flux.hardtanh.(x) + x = hardtanh.(x) return x ./ sum(abs.(x), dims=2) end From 15f449bb5353bf60705e9c77d49a24392aad410b Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Tue, 3 Mar 2026 23:13:31 -0500 Subject: [PATCH 076/120] up --- experiments/embeddings.jl | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 experiments/embeddings.jl diff --git a/experiments/embeddings.jl b/experiments/embeddings.jl new file mode 100644 index 0000000..2fc7fac --- /dev/null +++ b/experiments/embeddings.jl @@ -0,0 +1,36 @@ +using Random +using Lux +using DataFrames + +rng = Random.Xoshiro(123) +batch = 16 +nfeats = 3 +embed_lvls = 7 +embed_size = 9 + +struct NumCatEmbeddings{L} <: AbstractLuxWrapperLayer{:layer} + layer::L +end + +function NumCatEmbeddings(; num_size=5, cat_size=0, embed_lvls=7, embed_size=9) + + # numerical embeddings + nums = BatchNorm(num_size) + + # categorical embeddings + cats = [] + for i in 1:cat_size + push!(cats, Symbol("embed$i") => Embedding(embed_lvls => embed_size)) + end + m = Parallel(vcat; name="preproc", nums, cats...) + return NumCatEmbeddings(m) +end + +m = NumCatEmbeddings(; num_size=nfeats, cat_size=2) +ps, st = LuxCore.setup(rng, m) + +x = randn(rng, Float32, nfeats, batch) +c1 = rand(rng, 1:embed_lvls, batch) +c2 = rand(rng, 1:embed_lvls, batch) + +out, _st = m((x, c1, c2), ps, Lux.testmode(st)) From 9d3cb94b27e0ce1664dd68421e64a4724ada9db7 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 3 Mar 2026 23:23:53 -0500 Subject: [PATCH 077/120] up --- benchmarks/YEAR-regression.jl | 6 +++--- src/Fit/fit.jl | 2 +- src/models/NeuroTree/neurotrees.jl | 10 ++++++++-- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index ee8c312..7cc604a 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -65,13 +65,13 @@ arch = NeuroTabModels.NeuroTreeConfig(; # MLE_tree_split=false # ) -device = :cpu +device = :gpu loss = :mse # :mse :gaussian_mle :tweedie learner = NeuroTabRegressor( arch; loss, - nrounds=20, + nrounds=200, early_stopping_rounds=2, lr=1e-3, batchsize=1024, @@ -81,7 +81,7 @@ learner = NeuroTabRegressor( m = NeuroTabModels.fit( learner, dtrain; - # deval, + deval, target_name, feature_names, print_every_n=5, diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index cd39c2f..955d444 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -10,7 +10,7 @@ using ..Metrics import Random: Xoshiro import MLJModelInterface: fit -import Optimisers: OptimiserChain, WeightDecay, NAdam +import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam using Lux using Reactant diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 6ccef12..f400305 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -5,7 +5,7 @@ export NeuroTreeConfig using Random using Lux using LuxCore -using NNlib: softplus, sigmoid, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh +using NNlib: softplus, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh, tanhshrink import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture @@ -85,7 +85,7 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) layers = Any[create_tree(n_in, config.hidden_size)] - for _ in 1:(config.stack_size - 2) + for _ in 1:(config.stack_size-2) push!(layers, SkipConnection(create_tree(config.hidden_size, config.hidden_size), +)) end @@ -126,12 +126,17 @@ function _hardtanh_act(x) x = hardtanh.(x) return x ./ sum(abs.(x), dims=2) end +function _tanhshrink_act(x) + x = tanhshrink.(x) + return x ./ sum(abs.(x), dims=2) +end """ act_dict = Dict( :identity => _identity_act, :tanh => _tanh_act, :hardtanh => _hardtanh_act, + :tanhshrink => _tanhshrink_act, ) Dictionary mapping features activation name to their function. @@ -140,6 +145,7 @@ const act_dict = Dict( :identity => _identity_act, :tanh => _tanh_act, :hardtanh => _hardtanh_act, + :tanhshrink => _tanhshrink_act, ) end \ No newline at end of file From 2a169201a32ae4bc5da9cfe0a443b41c15baccfd Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Fri, 6 Mar 2026 12:30:21 +0530 Subject: [PATCH 078/120] Use AbstractLuxWrapperLayer for Neurotrees (#30) * Use AbstractLuxWrapperLayer to implement stacked NeuroTree Signed-off-by: AdityaPandeyCN * align changes Signed-off-by: AdityaPandeyCN * use shorthand named-tuple syntax for config kwargs Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN --- src/models/NeuroTree/neurotrees.jl | 67 ++++++++++++++++-------------- 1 file changed, 36 insertions(+), 31 deletions(-) diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 6ccef12..5f8453c 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -12,6 +12,26 @@ import ..Models: Architecture include("model.jl") +struct StackedNeuroTree{L} <: LuxCore.AbstractLuxWrapperLayer{:chain} + chain::L +end + +function StackedNeuroTree(; feats::Int, outs::Int, hidden_size::Int, stack_size::Int, tree_kwargs...) + if stack_size == 1 + return StackedNeuroTree(NeuroTree(; feats, outs, tree_kwargs...)) + end + + layers = Any[NeuroTree(; feats, outs=hidden_size, tree_kwargs...)] + for _ in 1:(stack_size - 2) + push!(layers, SkipConnection( + NeuroTree(; feats=hidden_size, outs=hidden_size, tree_kwargs...), + + )) + end + push!(layers, NeuroTree(; feats=hidden_size, outs, tree_kwargs...)) + + return StackedNeuroTree(Chain(layers...)) +end + struct NeuroTreeConfig <: Architecture tree_type::Symbol actA::Symbol @@ -66,33 +86,19 @@ function NeuroTreeConfig(; kwargs...) ) end -function (config::NeuroTreeConfig)(; nfeats, outsize) - function build_block(n_in, n_out) - create_tree(in_dim, out_dim) = NeuroTree(; - feats=in_dim, - outs=out_dim, - tree_type=config.tree_type, - depth=config.depth, - trees=config.ntrees, - actA=act_dict[config.actA], - scaler=config.scaler, - init_scale=config.init_scale - ) - - if config.stack_size == 1 - return create_tree(n_in, n_out) - end - - layers = Any[create_tree(n_in, config.hidden_size)] - - for _ in 1:(config.stack_size - 2) - push!(layers, SkipConnection(create_tree(config.hidden_size, config.hidden_size), +)) - end - - push!(layers, create_tree(config.hidden_size, n_out)) +function _tree_kwargs(config::NeuroTreeConfig) + return (; + config.tree_type, + config.depth, + trees = config.ntrees, + actA = act_dict[config.actA], + config.scaler, + config.init_scale, + ) +end - return Chain(layers...) - end +function (config::NeuroTreeConfig)(; nfeats, outsize) + kwargs = _tree_kwargs(config) if config.MLE_tree_split iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") @@ -101,14 +107,14 @@ function (config::NeuroTreeConfig)(; nfeats, outsize) BatchNorm(nfeats, track_stats=false), Parallel( vcat, - build_block(nfeats, head_outsize), - build_block(nfeats, head_outsize), - ) + StackedNeuroTree(; feats=nfeats, outs=head_outsize, config.hidden_size, config.stack_size, kwargs...), + StackedNeuroTree(; feats=nfeats, outs=head_outsize, config.hidden_size, config.stack_size, kwargs...), + ), ) else chain = Chain( BatchNorm(nfeats), - build_block(nfeats, outsize) + StackedNeuroTree(; feats=nfeats, outs=outsize, config.hidden_size, config.stack_size, kwargs...), ) end @@ -133,7 +139,6 @@ end :tanh => _tanh_act, :hardtanh => _hardtanh_act, ) - Dictionary mapping features activation name to their function. """ const act_dict = Dict( From abf767dd3dda90652ea15baabb131ab822cce387 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 10 Mar 2026 01:22:05 -0400 Subject: [PATCH 079/120] up --- benchmarks/YEAR-regression.jl | 2 +- benchmarks/titanic-mlogloss.jl | 25 ++++++++++++++++--------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 7cc604a..8ba487d 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -73,7 +73,7 @@ learner = NeuroTabRegressor( loss, nrounds=200, early_stopping_rounds=2, - lr=1e-3, + lr=3e-4, batchsize=1024, device ) diff --git a/benchmarks/titanic-mlogloss.jl b/benchmarks/titanic-mlogloss.jl index b1bf089..b5d32f7 100644 --- a/benchmarks/titanic-mlogloss.jl +++ b/benchmarks/titanic-mlogloss.jl @@ -34,24 +34,31 @@ deval = df[setdiff(1:nrow(df), train_indices), :] target_name = "y_cat" feature_names = setdiff(names(df), ["y_cat", "Survived"]) - eltype(dtrain[:, "y_cat"]) -config = NeuroTabClassifier( - nrounds=400, + +arch = NeuroTabModels.NeuroTreeConfig(; + actA=:identity, + init_scale=1.0, depth=4, - lr=3e-2, + ntrees=16, + stack_size=1, + hidden_size=1, +) + +learner = NeuroTabClassifier( + arch; + nrounds=100, + early_stopping_rounds=2, + lr=1e-2, ) m = NeuroTabModels.fit( - config, + learner, dtrain; deval, target_name, feature_names, - metric=:mlogloss, - print_every_n=10, - early_stopping_rounds=3, - device=:cpu + print_every_n=2, ) p_train = m(dtrain) From bf55f45965ff0d70f963053c7e8ea72c3c297639 Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Sun, 15 Mar 2026 00:45:06 +0530 Subject: [PATCH 080/120] Add documentation for the loss function (#32) * add docs Signed-off-by: AdityaPandeyCN * remove inferece reference Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN --- docs/src/losses.md | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 docs/src/losses.md diff --git a/docs/src/losses.md b/docs/src/losses.md new file mode 100644 index 0000000..faf71e2 --- /dev/null +++ b/docs/src/losses.md @@ -0,0 +1,38 @@ +# Losses + +## Prediction Shape + +All losses expect 3D predictions: `(outsize, K, batch)` where `K` is the ensemble size. 2D outputs are reshaped to `(outsize, 1, batch)` automatically. + +`reduce_pred` averages over `K` on raw predictions before any transformation. + +## Usage +```julia +get_loss_fn(:mse) # by symbol +get_loss_fn(MSE) # by type +get_loss_type(:mse) # → MSE +``` + +## Supported Losses + +| Symbol | Type | Pred shape | Target | Notes | +|--------|------|-----------|--------|-------| +| `:mse` | `MSE` | `(1, K, B)` | scalar | | +| `:mae` | `MAE` | `(1, K, B)` | scalar | | +| `:logloss` | `LogLoss` | `(1, K, B)` | `{0, 1}` | raw logits | +| `:mlogloss` | `MLogLoss` | `(C, K, B)` | `{1, …, C}` | raw logits | +| `:gaussian_mle` | `GaussianMLE` | `(2, K, B)` | scalar | `pred[1,:,:]` = μ, `pred[2,:,:]` = log-σ | +| `:tweedie` | `Tweedie` | `(1, K, B)` | non-negative | log-scale pred, ρ = 1.5 | + +## Data Tuples + +| Tuple | Contents | +|-------|----------| +| `(x, y)` | standard training | +| `(x, y, w)` | weighted training | +| `(x, y, w, offset)` | with offset (e.g. boosting) | + +## Signature +```julia +loss_fn(model, ps, st, data) → (scalar_loss, updated_state, NamedTuple()) +``` \ No newline at end of file From 6a763abb760acc0e319e8c07bc7c8ca4bf212301 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Wed, 18 Mar 2026 10:23:41 -0400 Subject: [PATCH 081/120] up --- .gitignore | 2 +- docs/Project.toml | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 08cfbbf..89f38b3 100644 --- a/.gitignore +++ b/.gitignore @@ -23,7 +23,7 @@ docs/.vscode docs/node_modules/ docs/.vitepress/cache docs/.vitepress/dist -# docs/package-lock.json +docs/package-lock.json # File generated by Pkg, the package manager, based on a corresponding Project.toml # It records a fixed state of all packages used by the project. As such, it should not be diff --git a/docs/Project.toml b/docs/Project.toml index 6cf1607..5ff744d 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,6 +1,7 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365" +NeuroTabModels = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" [compat] Documenter = "1" From 5cb5bd691e93ca87b0a80083c79cccf4286ef32a Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Wed, 18 Mar 2026 20:33:50 +0530 Subject: [PATCH 082/120] Add TabM architecture and numerical embeddings module (#29) * add numerical embeddings Signed-off-by: AdityaPandeyCN * add TabM architecture support Signed-off-by: AdityaPandeyCN * revert loss function signatures Signed-off-by: AdityaPandeyCN * handle ensemble Signed-off-by: AdityaPandeyCN * fix inference and embeddings Signed-off-by: AdityaPandeyCN * fix inference Signed-off-by: AdityaPandeyCN * fix embeddings Signed-off-by: AdityaPandeyCN * cleanup Signed-off-by: AdityaPandeyCN * broadcast ensemble losses over 3D predictions, drop 2D fallback Signed-off-by: AdityaPandeyCN * comments Signed-off-by: AdityaPandeyCN * losses fix Signed-off-by: AdityaPandeyCN * up * clean up TabM layers, unify loss handling for 2D/3D output Signed-off-by: AdityaPandeyCN * gaussian_mle loss fix Signed-off-by: AdityaPandeyCN * up * up * add tests for embedding module and TabM architecture Signed-off-by: AdityaPandeyCN * fix multi-class and refactor losses file Signed-off-by: AdityaPandeyCN * set compat for Reactant Signed-off-by: AdityaPandeyCN * fix version Signed-off-by: AdityaPandeyCN * restructure test Signed-off-by: AdityaPandeyCN * refactor infer.jl, deduplicate _broadcast_relu, remove unused LinearReLUEmbeddings Rename _activation to _inverse_link for proper GLM terminology. Apply inverse link after assembly instead of per-batch; add raw prediction option. Drop state return from _forward_reduce since inference only needs pred. Use pred/p naming consistently throughout infer.jl. Move reduce_pred from Losses to Infer, used by both Infer and Fit/callback. Define _broadcast_relu once in Models.jl, import in TabM and embeddings. Remove unused LinearReLUEmbeddings from linear.jl. Signed-off-by: AdityaPandeyCN * unify broadcast and ensure_3d to common reshape Signed-off-by: AdityaPandeyCN * use activation function instead of bool and clamp PLE to [0,1] Signed-off-by: AdityaPandeyCN * bump reactant version and use activation function in linear embedding Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN Co-authored-by: jeremie.desgagne.bouchard --- Project.toml | 2 + benchmarks/YEAR-regression.jl | 14 +- benchmarks/benchmark_mse.jl | 32 ++- benchmarks/titanic-logloss.jl | 38 ++-- benchmarks/titanic-mlogloss.jl | 25 ++- src/Fit/callback.jl | 7 +- src/Fit/fit.jl | 8 +- src/NeuroTabModels.jl | 7 +- src/infer.jl | 67 ++++--- src/losses.jl | 155 +++++---------- src/models/NeuroTree/neurotrees.jl | 6 +- src/models/TabM/TabM.jl | 206 +++++++++++++++++++ src/models/TabM/layers.jl | 229 ++++++++++++++++++++++ src/models/embeddings/compute_bins.jl | 48 +++++ src/models/embeddings/embeddings.jl | 21 ++ src/models/embeddings/linear.jl | 42 ++++ src/models/embeddings/nlinear.jl | 50 +++++ src/models/embeddings/periodic.jl | 95 +++++++++ src/models/embeddings/piecewise_linear.jl | 137 +++++++++++++ src/models/models.jl | 15 +- test/core.jl | 89 ++++++++- test/embedding.jl | 42 ++++ test/runtests.jl | 1 + 23 files changed, 1150 insertions(+), 186 deletions(-) create mode 100644 src/models/TabM/TabM.jl create mode 100644 src/models/TabM/layers.jl create mode 100644 src/models/embeddings/compute_bins.jl create mode 100644 src/models/embeddings/embeddings.jl create mode 100644 src/models/embeddings/linear.jl create mode 100644 src/models/embeddings/nlinear.jl create mode 100644 src/models/embeddings/periodic.jl create mode 100644 src/models/embeddings/piecewise_linear.jl create mode 100644 test/embedding.jl diff --git a/Project.toml b/Project.toml index 168673f..092d19c 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" +LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" @@ -29,6 +30,7 @@ MLUtils = "0.4" NNlib = "0.9" Optimisers = "0.4" Random = "1" +Reactant = "0.2.236" Statistics = "1" StatsBase = "0.34" Tables = "1.9" diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 7cc604a..cd83587 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -52,6 +52,18 @@ arch = NeuroTabModels.NeuroTreeConfig(; scaler=true, MLE_tree_split=false ) +# arch = NeuroTabModels.TabMConfig(; +# arch_type=:tabm, +# k=32, +# d_block=128, +# n_blocks=3, +# dropout=0.1, +# n_bins=16, +# use_embeddings=true, +# embedding_type=:linear, # periodic, piecewise, linear +# d_embedding=8, +# scaling_init=:normal, +# ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, # stack_size=1, @@ -66,7 +78,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # ) device = :gpu -loss = :mse # :mse :gaussian_mle :tweedie +loss = :gaussian_mle # :mse :gaussian_mle :tweedie learner = NeuroTabRegressor( arch; diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index b819565..f5a6b14 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -16,16 +16,28 @@ feature_names = names(dtrain) dtrain.y = Y target_name = "y" -arch = NeuroTabModels.NeuroTreeConfig(; - tree_type=:binary, - proj_size=1, - actA=:identity, - init_scale=1.0, - depth=4, - ntrees=32, - stack_size=1, - hidden_size=1, - scaler=false, +# arch = NeuroTabModels.NeuroTreeConfig(; +# tree_type=:binary, +# proj_size=1, +# actA=:identity, +# init_scale=1.0, +# depth=4, +# ntrees=32, +# stack_size=1, +# hidden_size=1, +# scaler=false, +# ) +arch = NeuroTabModels.TabMConfig(; + arch_type=:tabm, + k=16, + d_block=64, + n_blocks=3, + dropout=0.1, + bins=nothing, + use_embeddings=false, + embedding_type=:periodic, + d_embedding=16, + scaling_init=:random_signs, ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 8ef53de..4137b7a 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -32,31 +32,43 @@ deval = df[setdiff(1:nrow(df), train_indices), :] target_name = "Survived" feature_names = setdiff(names(df), ["Survived"]) -arch = NeuroTabModels.NeuroTreeConfig(; - tree_type=:binary, - proj_size=1, - init_scale=1.0, - depth=4, - ntrees=16, - stack_size=1, - hidden_size=1, - actA=:identity, - MLE_tree_split=false, +# arch = NeuroTabModels.NeuroTreeConfig(; +# tree_type=:binary, +# proj_size=1, +# init_scale=1.0, +# depth=4, +# ntrees=16, +# stack_size=1, +# hidden_size=1, +# actA=:identity, +# MLE_tree_split=false, +# ) +arch = NeuroTabModels.TabMConfig(; + arch_type=:tabm, + k=32, + d_block=64, + n_blocks=3, + dropout=0.1, + bins=nothing, + use_embeddings=false, + embedding_type=:periodic, + d_embedding=16, + scaling_init=:random_signs, ) learner = NeuroTabRegressor( arch; - loss=:logloss, # FIXME: gaussian_mle don't train + loss=:logloss, nrounds=100, early_stopping_rounds=2, - lr=3e-2, + lr=1e-2, device=:cpu ) @time m = NeuroTabModels.fit( learner, dtrain; - # deval, # FIXME: important slowdown when deval is used + deval, target_name, feature_names, print_every_n=10, diff --git a/benchmarks/titanic-mlogloss.jl b/benchmarks/titanic-mlogloss.jl index b1bf089..b5d32f7 100644 --- a/benchmarks/titanic-mlogloss.jl +++ b/benchmarks/titanic-mlogloss.jl @@ -34,24 +34,31 @@ deval = df[setdiff(1:nrow(df), train_indices), :] target_name = "y_cat" feature_names = setdiff(names(df), ["y_cat", "Survived"]) - eltype(dtrain[:, "y_cat"]) -config = NeuroTabClassifier( - nrounds=400, + +arch = NeuroTabModels.NeuroTreeConfig(; + actA=:identity, + init_scale=1.0, depth=4, - lr=3e-2, + ntrees=16, + stack_size=1, + hidden_size=1, +) + +learner = NeuroTabClassifier( + arch; + nrounds=100, + early_stopping_rounds=2, + lr=1e-2, ) m = NeuroTabModels.fit( - config, + learner, dtrain; deval, target_name, feature_names, - metric=:mlogloss, - print_every_n=10, - early_stopping_rounds=3, - device=:cpu + print_every_n=2, ) p_train = m(dtrain) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 423c4b4..8ac5f76 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -4,6 +4,7 @@ using DataFrames using Statistics: mean, median using ..Learners: LearnerTypes +using ..Infer: reduce_pred using ..Data: get_df_loader_train using ..Metrics @@ -47,19 +48,19 @@ end function _compile_eval_step(chain, feval, d0, ps, st) if length(d0) == 2 function _step2(x, y, ps, st) - m = x -> first(chain(x, ps, st)) + m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y; agg=sum), eltype(y)(last(size(y))) end return @compile _step2(d0[1], d0[2], ps, st) elseif length(d0) == 3 function _step3(x, y, w, ps, st) - m = x -> first(chain(x, ps, st)) + m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y, w; agg=sum), sum(w) end return @compile _step3(d0[1], d0[2], d0[3], ps, st) else function _step4(x, y, w, offset, ps, st) - m = x -> first(chain(x, ps, st)) + m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y, w, offset; agg=sum), sum(w) end return @compile _step4(d0[1], d0[2], d0[3], d0[4], ps, st) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 955d444..da2ece6 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -7,6 +7,7 @@ using ..Learners using ..Models using ..Losses using ..Metrics +using ..Infer import Random: Xoshiro import MLJModelInterface: fit @@ -65,7 +66,12 @@ function init( :device => config.device ) - chain = config.arch(; nfeats, outsize) + if hasproperty(config.arch, :use_embeddings) && config.arch.use_embeddings && config.arch.embedding_type == :piecewise + X_train = Matrix{Float32}(df[:, feature_names]) + chain = config.arch(; nfeats, outsize, X_train) + else + chain = config.arch(; nfeats, outsize) + end m = NeuroTabModel(L, chain, info) rng = Xoshiro(config.seed) diff --git a/src/NeuroTabModels.jl b/src/NeuroTabModels.jl index ccabb56..a13c884 100644 --- a/src/NeuroTabModels.jl +++ b/src/NeuroTabModels.jl @@ -3,18 +3,19 @@ module NeuroTabModels using Random export NeuroTabRegressor, NeuroTabClassifier, NeuroTabModel +export TabMConfig, compute_bins include("data.jl") include("losses.jl") include("metrics.jl") include("models/models.jl") using .Models -include("learners.jl") -using .Learners include("infer.jl") using .Infer +include("learners.jl") +using .Learners include("Fit/fit.jl") using .Fit include("MLJ.jl") -end # module +end # module \ No newline at end of file diff --git a/src/infer.jl b/src/infer.jl index ff27c94..f1b9745 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -3,17 +3,20 @@ module Infer using ..Data using ..Losses using ..Models -using ..Learners using Lux using Lux: cpu_device, reactant_device using Reactant using Reactant: @compile using NNlib: sigmoid, softmax +using Statistics: mean using DataFrames: AbstractDataFrame import MLUtils: DataLoader -export infer +export infer, reduce_pred + +reduce_pred(pred::AbstractMatrix) = pred +reduce_pred(pred::AbstractArray{T,3}) where {T} = dropdims(mean(pred; dims=2); dims=2) function _get_device(device::Symbol) backend = device == :gpu ? "gpu" : "cpu" @@ -21,34 +24,38 @@ function _get_device(device::Symbol) return reactant_device() end -function _postprocess(::Type{<:Union{MSE,MAE}}, raw_preds) - return vcat([vec(p) for p in raw_preds]...) +function _forward_reduce(chain, x, ps, st) + pred, _ = chain(x, ps, st) + return reduce_pred(pred) end -function _postprocess(::Type{<:LogLoss}, raw_preds) - p = vcat([vec(p) for p in raw_preds]...) - return sigmoid.(p) -end +# Assemble raw predictions into final structure (no transforms) +_assemble(::Type{<:MLogLoss}, raw_preds) = reduce(hcat, raw_preds) +_assemble(::Type{<:GaussianMLE}, raw_preds) = reduce(hcat, raw_preds) +_assemble(::Type, raw_preds) = vcat([vec(p) for p in raw_preds]...) + +# Apply inverse link to convert from model scale to natural scale +_inverse_link(::Type{<:LogLoss}, pred) = sigmoid.(pred) +_inverse_link(::Type{<:Tweedie}, pred) = exp.(pred) +_inverse_link(::Type{<:Union{MSE,MAE}}, pred) = pred -function _postprocess(::Type{<:MLogLoss}, raw_preds) - p_full = reduce(hcat, raw_preds) - p_soft = softmax(p_full; dims=1) - return Matrix(p_soft') +function _inverse_link(::Type{<:MLogLoss}, pred) + return Matrix(softmax(pred; dims=1)') end -function _postprocess(::Type{<:GaussianMLE}, raw_preds) - p_full = reduce(hcat, raw_preds) - p_T = Matrix(p_full') - p_T[:, 2] .= exp.(p_T[:, 2]) - return p_T +function _inverse_link(::Type{<:GaussianMLE}, pred) + p = Matrix(pred') + p[:, 2] .= exp.(p[:, 2]) + return p end -function _postprocess(::Type{<:Tweedie}, raw_preds) - p = vcat([vec(p) for p in raw_preds]...) - return exp.(p) +function _postprocess(::Type{L}, raw_preds; raw::Bool=false) where {L} + assembled = _assemble(L, raw_preds) + raw && return assembled + return _inverse_link(L, assembled) end -function infer(m::NeuroTabModel{L}, data; device=:cpu) where {L} +function infer(m::NeuroTabModel{L}, data; device=:cpu, raw::Bool=false) where {L} dev = _get_device(device) cdev = cpu_device() ps = dev(m.info[:ps]) @@ -58,32 +65,32 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu) where {L} b_first = first(data) x0 = b_first isa Tuple ? b_first[1] : b_first - model_compiled = @compile m.chain(dev(x0), ps, st) + compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) for b in data x = b isa Tuple ? b[1] : b if size(x) == size(x0) - y_pred, _ = model_compiled(dev(x), ps, st) + pred = compiled(m.chain, dev(x), ps, st) else - y_pred, _ = Reactant.@jit Lux.apply(m.chain, dev(x), ps, st) + pred = Reactant.@jit _forward_reduce(m.chain, dev(x), ps, st) end - push!(raw_preds, cdev(y_pred)) + push!(raw_preds, cdev(pred)) end - return _postprocess(L, raw_preds) + return _postprocess(L, raw_preds; raw) end -function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu) +function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu, raw::Bool=false) dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) - return infer(m, dinfer; device=device) + return infer(m, dinfer; device, raw) end function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) - return infer(m, data; device=device) + return infer(m, data; device) end function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) - return infer(m, [(x,)]; device=device) + return infer(m, [(x,)]; device) end end \ No newline at end of file diff --git a/src/losses.jl b/src/losses.jl index 920b602..2648d47 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -14,123 +14,72 @@ abstract type MLogLoss <: LossType end abstract type GaussianMLE <: LossType end abstract type Tweedie <: LossType end -function mse_loss(model, ps, st, data::Tuple{Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred); y = vec(data[2]) - return mean((p .- y) .^ 2), st_, NamedTuple() -end -function mse_loss(model, ps, st, data::Tuple{Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred); y = vec(data[2]); w = vec(data[3]) - return sum((p .- y) .^ 2 .* w) / sum(w), st_, NamedTuple() -end -function mse_loss(model, ps, st, data::Tuple{Any,Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred) .+ vec(data[4]); y = vec(data[2]); w = vec(data[3]) - return sum((p .- y) .^ 2 .* w) / sum(w), st_, NamedTuple() -end +_reshape_3d(x::AbstractVector) = reshape(x, 1, 1, :) +_reshape_3d(x::AbstractMatrix) = reshape(x, size(x, 1), 1, size(x, 2)) +_reshape_3d(x::AbstractArray{T,3}) where {T} = x -function mae_loss(model, ps, st, data::Tuple{Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred); y = vec(data[2]) - return mean(abs.(p .- y)), st_, NamedTuple() -end -function mae_loss(model, ps, st, data::Tuple{Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred); y = vec(data[2]); w = vec(data[3]) - return sum(abs.(p .- y) .* w) / sum(w), st_, NamedTuple() -end -function mae_loss(model, ps, st, data::Tuple{Any,Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred) .+ vec(data[4]); y = vec(data[2]); w = vec(data[3]) - return sum(abs.(p .- y) .* w) / sum(w), st_, NamedTuple() +function _forward(model, ps, st, x) + pred, st_ = model(x, ps, st) + return _reshape_3d(pred), st_ end -function logloss(model, ps, st, data::Tuple{Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred); y = vec(data[2]) - return mean((1 .- y) .* p .- logsigmoid.(p)), st_, NamedTuple() -end -function logloss(model, ps, st, data::Tuple{Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred); y = vec(data[2]); w = vec(data[3]) - return sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w), st_, NamedTuple() -end -function logloss(model, ps, st, data::Tuple{Any,Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - p = vec(pred) .+ vec(data[4]); y = vec(data[2]); w = vec(data[3]) - return sum(w .* ((1 .- y) .* p .- logsigmoid.(p))) / sum(w), st_, NamedTuple() -end +_reduce(loss) = mean(loss) +_reduce(loss, w) = sum(mean(loss; dims=2) .* w) / sum(w) -function mlogloss(model, ps, st, data::Tuple{Any,Any}) - pred, st_ = model(data[1], ps, st) - k = size(pred, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(data[2], 1, :) - lsm = logsoftmax(pred; dims=1) - return mean(-sum(y_oh .* lsm; dims=1)), st_, NamedTuple() +function _apply_loss(core, model, ps, st, data::Tuple{Any,Any}) + pred, st_ = _forward(model, ps, st, data[1]) + return _reduce(core(pred, _reshape_3d(data[2]))), st_, NamedTuple() end -function mlogloss(model, ps, st, data::Tuple{Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - k = size(pred, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(data[2], 1, :) - lsm = logsoftmax(pred; dims=1) - return sum(vec(-sum(y_oh .* lsm; dims=1)) .* vec(data[3])) / sum(data[3]), st_, NamedTuple() +function _apply_loss(core, model, ps, st, data::Tuple{Any,Any,Any}) + pred, st_ = _forward(model, ps, st, data[1]) + return _reduce(core(pred, _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() end -function mlogloss(model, ps, st, data::Tuple{Any,Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - pred = pred .+ data[4] - k = size(pred, 1) - y_oh = (UInt32(1):UInt32(k)) .== reshape(data[2], 1, :) - lsm = logsoftmax(pred; dims=1) - return sum(vec(-sum(y_oh .* lsm; dims=1)) .* vec(data[3])) / sum(data[3]), st_, NamedTuple() +function _apply_loss(core, model, ps, st, data::Tuple{Any,Any,Any,Any}) + pred, st_ = _forward(model, ps, st, data[1]) + return _reduce(core(pred .+ _reshape_3d(data[4]), _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() end -function tweedie(model, ps, st, data::Tuple{Any,Any}) - pred, st_ = model(data[1], ps, st) - rho = eltype(data[1])(1.5) - ep = exp.(vec(pred)); y = vec(data[2]) - loss = mean(2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ - ep .^ (2 - rho) / (2 - rho))) - return loss, st_, NamedTuple() -end -function tweedie(model, ps, st, data::Tuple{Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - rho = eltype(data[1])(1.5) - ep = exp.(vec(pred)); y = vec(data[2]); w = vec(data[3]) - loss = sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ - ep .^ (2 - rho) / (2 - rho))) / sum(w) - return loss, st_, NamedTuple() +_mse_core(pred, y) = (pred .- y) .^ 2 +_mae_core(pred, y) = abs.(pred .- y) +_logloss_core(pred, y) = (1 .- y) .* pred .- logsigmoid.(pred) + +function _mlogloss_core(pred, y) + nclasses = size(pred, 1) + classes = reshape(Int32(1):Int32(nclasses), :, 1, 1) + y_idx = reshape(Int32.(y), 1, 1, :) + y_oh = Float32.(classes .== y_idx) + return -sum(y_oh .* logsoftmax(pred; dims=1); dims=1) end -function tweedie(model, ps, st, data::Tuple{Any,Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - rho = eltype(data[1])(1.5) - ep = exp.(vec(pred) .+ vec(data[4])); y = vec(data[2]); w = vec(data[3]) - loss = sum(w .* 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ - ep .^ (2 - rho) / (2 - rho))) / sum(w) - return loss, st_, NamedTuple() + +function _tweedie_core(pred, y) + rho = eltype(pred)(1.5) + ep = exp.(pred) + 2 .* (y .^ (2 - rho) / (1 - rho) / (2 - rho) .- y .* ep .^ (1 - rho) / (1 - rho) .+ + ep .^ (2 - rho) / (2 - rho)) end -gaussian_mle_loss(μ::AbstractVector, σ::AbstractVector, y::AbstractVector) = - -sum(-σ .- (y .- μ) .^ 2 ./ (2 .* max.(oftype.(σ, 2e-7), exp.(2 .* σ)))) +mse_loss(m, ps, st, d) = _apply_loss(_mse_core, m, ps, st, d) +mae_loss(m, ps, st, d) = _apply_loss(_mae_core, m, ps, st, d) +logloss(m, ps, st, d) = _apply_loss(_logloss_core, m, ps, st, d) +mlogloss(m, ps, st, d) = _apply_loss(_mlogloss_core, m, ps, st, d) +tweedie(m, ps, st, d) = _apply_loss(_tweedie_core, m, ps, st, d) -gaussian_mle_loss(μ::AbstractVector, σ::AbstractVector, y::AbstractVector, w::AbstractVector) = - -sum((-σ .- (y .- μ) .^ 2 ./ (2 .* max.(oftype.(σ, 2e-7), exp.(2 .* σ)))) .* w) / sum(w) +function _gaussian_mle_core(μ, σ, y) + σ .+ (y .- μ) .^ 2 ./ (2 .* max.(eltype(σ)(2e-7), exp.(2 .* σ))) +end function gaussian_mle(model, ps, st, data::Tuple{Any,Any}) - pred, st_ = model(data[1], ps, st) - loss = gaussian_mle_loss(view(pred, 1, :), view(pred, 2, :), vec(data[2])) - return loss, st_, NamedTuple() + pred, st_ = _forward(model, ps, st, data[1]) + return _reduce(_gaussian_mle_core(pred[1:1,:,:], pred[2:2,:,:], _reshape_3d(data[2]))), st_, NamedTuple() end function gaussian_mle(model, ps, st, data::Tuple{Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - loss = gaussian_mle_loss(view(pred, 1, :), view(pred, 2, :), vec(data[2]), vec(data[3])) - return loss, st_, NamedTuple() + pred, st_ = _forward(model, ps, st, data[1]) + return _reduce(_gaussian_mle_core(pred[1:1,:,:], pred[2:2,:,:], _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() end function gaussian_mle(model, ps, st, data::Tuple{Any,Any,Any,Any}) - pred, st_ = model(data[1], ps, st) - pred = pred .+ data[4] - loss = gaussian_mle_loss(view(pred, 1, :), view(pred, 2, :), vec(data[2]), vec(data[3])) - return loss, st_, NamedTuple() + pred, st_ = _forward(model, ps, st, data[1]) + pred = pred .+ _reshape_3d(data[4]) + return _reduce(_gaussian_mle_core(pred[1:1,:,:], pred[2:2,:,:], _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() end get_loss_fn(::Type{<:MSE}) = mse_loss @@ -141,12 +90,8 @@ get_loss_fn(::Type{<:GaussianMLE}) = gaussian_mle get_loss_fn(::Type{<:Tweedie}) = tweedie const _loss_type_dict = Dict( - :mse => MSE, - :mae => MAE, - :logloss => LogLoss, - :tweedie => Tweedie, - :gaussian_mle => GaussianMLE, - :mlogloss => MLogLoss, + :mse => MSE, :mae => MAE, :logloss => LogLoss, + :mlogloss => MLogLoss, :gaussian_mle => GaussianMLE, :tweedie => Tweedie, ) get_loss_type(loss::Symbol) = _loss_type_dict[loss] diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 51bc30c..42860bb 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -22,7 +22,7 @@ function StackedNeuroTree(; feats::Int, outs::Int, hidden_size::Int, stack_size: end layers = Any[NeuroTree(; feats, outs=hidden_size, tree_kwargs...)] - for _ in 1:(stack_size - 2) + for _ in 1:(stack_size-2) push!(layers, SkipConnection( NeuroTree(; feats=hidden_size, outs=hidden_size, tree_kwargs...), + )) @@ -90,8 +90,8 @@ function _tree_kwargs(config::NeuroTreeConfig) return (; config.tree_type, config.depth, - trees = config.ntrees, - actA = act_dict[config.actA], + trees=config.ntrees, + actA=act_dict[config.actA], config.scaler, config.init_scale, ) diff --git a/src/models/TabM/TabM.jl b/src/models/TabM/TabM.jl new file mode 100644 index 0000000..df4b386 --- /dev/null +++ b/src/models/TabM/TabM.jl @@ -0,0 +1,206 @@ +module TabM + +export TabMConfig + +using Random +using Lux +using LuxCore +using NNlib + +import ..Losses: get_loss_type, GaussianMLE +import ..Models: Architecture +import ..Embeddings: PeriodicEmbeddings, LinearEmbeddings, PiecewiseLinearEmbeddings, compute_bins +import ..Models: _broadcast_relu + +include("layers.jl") + +function _batch_ensemble_backbone(; + d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, + k::Int, scaling_init::Symbol, d_features::Vector{Int}) + layers = Any[] + for i in 1:n_blocks + d_in_i = (i == 1) ? d_in : d_block + if i == 1 + push!(layers, LinearBatchEnsemble(d_in_i, d_block; + k, scaling_init = (scaling_init, :ones), + first_scaling_init_chunks = d_features)) + else + push!(layers, LinearBatchEnsemble(d_in_i, d_block; + k, scaling_init = :ones)) + end + push!(layers, WrappedFunction(_broadcast_relu)) + dropout > 0 && push!(layers, Dropout(dropout)) + end + return layers +end + +function _mini_ensemble_backbone(; + d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, + k::Int, scaling_init::Symbol, d_features::Vector{Int}) + layers = Any[ScaleEnsemble(k, d_in; + init = scaling_init, init_chunks = d_features, bias = false)] + for i in 1:n_blocks + d_in_i = (i == 1) ? d_in : d_block + push!(layers, Dense(d_in_i => d_block, relu)) + dropout > 0 && push!(layers, Dropout(dropout)) + end + return layers +end + +function _packed_ensemble_backbone(; + d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, k::Int) + layers = Any[] + for i in 1:n_blocks + d_in_i = (i == 1) ? d_in : d_block + push!(layers, LinearEnsemble(d_in_i, d_block, k)) + push!(layers, WrappedFunction(_broadcast_relu)) + dropout > 0 && push!(layers, Dropout(dropout)) + end + return layers +end + +""" + TabMConfig(; kwargs...) + +Configuration for TabM ensemble architectures (Gorishniy et al., ICLR 2025). + +This implements the TabM♠ variant (shared training batches), where all k ensemble +members see the same training batch. + +The chain output is 3D: `(outsize, k, batch)`. Ensemble averaging is handled by +`Losses.reduce_pred` during training and `_reduce` in `infer.jl` at inference. + +# Arguments +- `k::Int`: Number of ensemble members (default `32`). +- `n_blocks::Int`: Number of MLP blocks (default `2` with embeddings, `3` without). +- `d_block::Int`: Hidden dimension per block (default `512`). +- `dropout::Float64`: Dropout rate (default `0.1`). +- `arch_type::Symbol`: `:tabm`, `:tabm_mini`, or `:tabm_packed` (default `:tabm`). +- `scaling_init::Symbol`: Init for ensemble scaling vectors — `:random_signs`, `:normal`, or `:ones` (default `:normal` with embeddings, `:random_signs` without). +- `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). +- `use_embeddings::Bool`: Apply feature embeddings before the backbone (default `false`). +- `d_embedding::Int`: Embedding dimension per feature (default `24`). +- `embedding_type::Symbol`: `:periodic`, `:linear`, or `:piecewise` (default `:periodic`). +- `n_bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). + A single `Int` applies the same count to all features. A `Vector{Int}` specifies + per-feature bin counts. + +# Callable + config(; nfeats, outsize, X_train=nothing) → Lux.Chain + +- `nfeats::Int`: Number of input features. +- `outsize::Int`: Output dimension. +- `X_train::Union{Nothing, AbstractMatrix}`: Training data of shape `(n_samples, n_features)`. + Required when `embedding_type=:piecewise` to compute bin edges via `compute_bins`. +""" +struct TabMConfig <: Architecture + k::Int + n_blocks::Int + d_block::Int + dropout::Float64 + arch_type::Symbol + scaling_init::Symbol + MLE_tree_split::Bool + use_embeddings::Bool + d_embedding::Int + embedding_type::Symbol + n_bins::Union{Int, Vector{Int}} +end + +function TabMConfig(; kwargs...) + has_embeddings = get(kwargs, :use_embeddings, false) + + args = Dict{Symbol,Any}( + :k => 32, + :n_blocks => has_embeddings ? 2 : 3, + :d_block => 512, + :dropout => 0.1, + :arch_type => :tabm, + :scaling_init => has_embeddings ? :normal : :random_signs, + :MLE_tree_split => false, + :use_embeddings => false, + :d_embedding => 24, + :embedding_type => :periodic, + :n_bins => 48, + ) + + args_ignored = setdiff(keys(kwargs), keys(args)) + length(args_ignored) > 0 && + @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(join(args_ignored, ", "))." + + args_default = setdiff(keys(args), keys(kwargs)) + length(args_default) > 0 && + @info "Following $(length(args_default)) arguments set to default: $(join(args_default, ", "))." + + for arg in intersect(keys(args), keys(kwargs)) + args[arg] = kwargs[arg] + end + + return TabMConfig( + args[:k], args[:n_blocks], args[:d_block], args[:dropout], + Symbol(args[:arch_type]), Symbol(args[:scaling_init]), + args[:MLE_tree_split], args[:use_embeddings], + args[:d_embedding], Symbol(args[:embedding_type]), args[:n_bins], + ) +end + +function (config::TabMConfig)(; nfeats, outsize, X_train=nothing) + @assert config.k > 0 "k must be > 0, got $(config.k)" + @assert nfeats > 0 "nfeats must be > 0, got $nfeats" + @assert outsize > 0 "outsize must be > 0, got $outsize" + + k = config.k + d_block = config.d_block + + if config.use_embeddings + emb_layer = if config.embedding_type == :periodic + PeriodicEmbeddings(nfeats, config.d_embedding) + elseif config.embedding_type == :linear + LinearEmbeddings(nfeats, config.d_embedding) + elseif config.embedding_type == :piecewise + @assert X_train !== nothing "Piecewise embeddings require `X_train` to compute bin edges." + bins = compute_bins(X_train; n_bins=config.n_bins) + @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))" + PiecewiseLinearEmbeddings(bins, config.d_embedding) + else + error("Unsupported embedding type: $(config.embedding_type)") + end + feature_layers = [emb_layer, FlattenLayer()] + d_in = nfeats * config.d_embedding + d_features = fill(config.d_embedding, nfeats) + effective_scaling_init = :normal + else + feature_layers = [] + d_in = nfeats + d_features = ones(Int, nfeats) + effective_scaling_init = config.scaling_init + end + + bb = if config.arch_type == :tabm + _batch_ensemble_backbone(; d_in, n_blocks=config.n_blocks, + d_block, dropout=config.dropout, k, + scaling_init=effective_scaling_init, d_features) + elseif config.arch_type == :tabm_mini + _mini_ensemble_backbone(; d_in, n_blocks=config.n_blocks, + d_block, dropout=config.dropout, k, + scaling_init=effective_scaling_init, d_features) + elseif config.arch_type == :tabm_packed + _packed_ensemble_backbone(; d_in, n_blocks=config.n_blocks, + d_block, dropout=config.dropout, k) + else + error("Unknown arch_type: $(config.arch_type)") + end + + head = if config.MLE_tree_split && outsize == 2 + split_out = outsize ÷ 2 + Parallel(vcat, + LinearEnsemble(d_block, split_out, k), + LinearEnsemble(d_block, split_out, k)) + else + LinearEnsemble(d_block, outsize, k) + end + + return Chain(feature_layers..., EnsembleView(k), bb..., head) +end + +end \ No newline at end of file diff --git a/src/models/TabM/layers.jl b/src/models/TabM/layers.jl new file mode 100644 index 0000000..695704a --- /dev/null +++ b/src/models/TabM/layers.jl @@ -0,0 +1,229 @@ +using Lux: AbstractLuxLayer, WrappedFunction +using LuxCore +using LuxLib: batched_matmul +using Random: AbstractRNG + +""" + _init_rsqrt_uniform(rng, dims, d) → Array{Float32} + +Uniform init in `[-1/√d, 1/√d]`. +""" +function _init_rsqrt_uniform(rng::AbstractRNG, dims, d::Int) + s = Float32(1 / sqrt(d)) + return s .* (2f0 .* rand(rng, Float32, dims...) .- 1f0) +end + +""" + _init_scaling(rng, dims, init) → Array{Float32} + +Initialize scaling vectors for ensemble adapters. +- `:ones` — deterministic ones +- `:normal` — N(0,1) +- `:random_signs` — ±1 +""" +function _init_scaling(rng::AbstractRNG, dims, init::Symbol) + if init == :ones + return ones(Float32, dims...) + elseif init == :normal + return randn(rng, Float32, dims...) + elseif init == :random_signs + return Float32.(2 .* (rand(rng, Float32, dims...) .> 0.5f0) .- 1) + else + error("Unknown scaling init: $init") + end +end + +""" + _init_scaling_with_chunks(rng, dims, init, chunks) → Array{Float32} + +Initialize scaling with grouped chunks. Each chunk shares the same random value +per ensemble member, providing structured diversity for features with different +representation sizes. +""" +function _init_scaling_with_chunks(rng::AbstractRNG, dims::Tuple{Int,Int}, + init::Symbol, chunks::Vector{Int}) + d, k = dims + @assert d == sum(chunks) "Chunks must sum to $d, got $(sum(chunks))" + weight = zeros(Float32, d, k) + row = 1 + for chunk_size in chunks + val = _init_scaling(rng, (1, k), init) + weight[row:row+chunk_size-1, :] .= repeat(val, chunk_size, 1) + row += chunk_size + end + return weight +end + +""" + EnsembleView(k) + +Broadcasts a `(D, B)` matrix to `(D, K, B)` by repeating along dim 2. +Passes through `(D, K, B)` input unchanged. +""" +struct EnsembleView <: AbstractLuxLayer + k::Int +end + +function (m::EnsembleView)(x::AbstractMatrix, ps, st) + D, B = size(x) + return repeat(reshape(x, D, 1, B), 1, m.k, 1), st +end + +function (m::EnsembleView)(x::AbstractArray{T,3}, ps, st) where {T} + @assert size(x, 2) == m.k "Expected K=$(m.k), got $(size(x, 2))" + return x, st +end + +""" + LinearBatchEnsemble(in_f, out_f; k, scaling_init=:random_signs, + first_scaling_init_chunks=nothing, bias=true) + +Batch-ensemble linear: `y = S ⊙ (W(R ⊙ x)) + bias` with shared `W` and +per-member `R`, `S`, `bias`. + +Equivalent to defining per-member weight matrices `Wᵢ = W ⊙ (sᵢrᵢᵀ)`. + +# Arguments +- `in_f::Int`: Input dimension. +- `out_f::Int`: Output dimension. +- `k::Int`: Number of ensemble members. +- `scaling_init`: Init for R/S. A `Symbol` (applied to both) or + `Tuple{Symbol,Symbol}` for `(R, S)`. Options: `:ones`, `:normal`, `:random_signs`. +- `first_scaling_init_chunks`: Chunk sizes for grouped R init. +- `bias::Bool`: Include per-member bias (default `true`). +""" +struct LinearBatchEnsemble <: AbstractLuxLayer + in_features::Int + out_features::Int + k::Int + use_bias::Bool + r_init::Symbol + s_init::Symbol + r_init_chunks::Union{Nothing, Vector{Int}} +end + +function LinearBatchEnsemble(in_f::Int, out_f::Int; + k::Int, + scaling_init::Union{Symbol, Tuple{Symbol,Symbol}} = :random_signs, + first_scaling_init_chunks::Union{Nothing, Vector{Int}} = nothing, + bias::Bool = true) + r_init, s_init = scaling_init isa Tuple ? scaling_init : (scaling_init, scaling_init) + return LinearBatchEnsemble(in_f, out_f, k, bias, r_init, s_init, first_scaling_init_chunks) +end + +function LuxCore.initialparameters(rng::AbstractRNG, m::LinearBatchEnsemble) + weight = _init_rsqrt_uniform(rng, (m.out_features, m.in_features), m.in_features) + + r = if m.r_init_chunks !== nothing + _init_scaling_with_chunks(rng, (m.in_features, m.k), m.r_init, m.r_init_chunks) + else + _init_scaling(rng, (m.in_features, m.k), m.r_init) + end + s = _init_scaling(rng, (m.out_features, m.k), m.s_init) + + d = (; weight, r, s) + if m.use_bias + bias = _init_rsqrt_uniform(rng, (m.out_features, m.k), m.in_features) + d = merge(d, (; bias)) + end + return d +end + +function (m::LinearBatchEnsemble)(x::AbstractArray{T,3}, ps, st) where {T} + in_f, k, batch = size(x) + x = x .* reshape(ps.r, m.in_features, m.k, 1) + x = reshape(ps.weight * reshape(x, in_f, k * batch), m.out_features, k, batch) + x = x .* reshape(ps.s, m.out_features, m.k, 1) + if m.use_bias + x = x .+ reshape(ps.bias, m.out_features, m.k, 1) + end + return x, st +end + +""" + LinearEnsemble(in_f, out_f, k; bias=true) + +`k` independent linear layers applied via `batched_matmul`. + +# Arguments +- `in_f::Int`: Input dimension. +- `out_f::Int`: Output dimension. +- `k::Int`: Number of ensemble members. +- `bias::Bool`: Include per-member bias (default `true`). +""" +struct LinearEnsemble <: AbstractLuxLayer + in_features::Int + out_features::Int + k::Int + use_bias::Bool +end + +LinearEnsemble(in_f::Int, out_f::Int, k::Int; bias::Bool = true) = + LinearEnsemble(in_f, out_f, k, bias) + +function LuxCore.initialparameters(rng::AbstractRNG, m::LinearEnsemble) + d = (; weight = _init_rsqrt_uniform(rng, (m.out_features, m.in_features, m.k), m.in_features)) + if m.use_bias + d = merge(d, (; bias = _init_rsqrt_uniform(rng, (m.out_features, m.k), m.in_features))) + end + return d +end + +function (m::LinearEnsemble)(x::AbstractArray{T,3}, ps, st) where {T} + xp = permutedims(x, (1, 3, 2)) + out = batched_matmul(ps.weight, xp) + out = permutedims(out, (1, 3, 2)) + if m.use_bias + out = out .+ reshape(ps.bias, m.out_features, m.k, 1) + end + return out, st +end + +""" + ScaleEnsemble(k, d; init=:random_signs, init_chunks=nothing, bias=false) + +Per-member elementwise scaling (and optional bias). + +# Arguments +- `k::Int`: Number of ensemble members. +- `d::Int`: Feature dimension. +- `init::Symbol`: Weight init — `:ones`, `:normal`, or `:random_signs`. +- `init_chunks`: Chunk sizes for grouped init. +- `bias::Bool`: Include per-member additive bias (default `false`). +""" +struct ScaleEnsemble <: AbstractLuxLayer + k::Int + d::Int + init::Symbol + init_chunks::Union{Nothing, Vector{Int}} + use_bias::Bool +end + +function ScaleEnsemble(k::Int, d::Int; + init::Symbol = :random_signs, + init_chunks::Union{Nothing, Vector{Int}} = nothing, + bias::Bool = false) + return ScaleEnsemble(k, d, init, init_chunks, bias) +end + +function LuxCore.initialparameters(rng::AbstractRNG, m::ScaleEnsemble) + weight = if m.init_chunks !== nothing + _init_scaling_with_chunks(rng, (m.d, m.k), m.init, m.init_chunks) + else + _init_scaling(rng, (m.d, m.k), m.init) + end + d = (; weight) + if m.use_bias + d = merge(d, (; bias = zeros(Float32, m.d, m.k))) + end + return d +end + +function (m::ScaleEnsemble)(x::AbstractArray{T,3}, ps, st) where {T} + w = reshape(ps.weight, m.d, m.k, 1) + if m.use_bias + return reshape(ps.bias, m.d, m.k, 1) .+ w .* x, st + else + return x .* w, st + end +end \ No newline at end of file diff --git a/src/models/embeddings/compute_bins.jl b/src/models/embeddings/compute_bins.jl new file mode 100644 index 0000000..1643379 --- /dev/null +++ b/src/models/embeddings/compute_bins.jl @@ -0,0 +1,48 @@ +""" + compute_bins(X; n_bins=48) + +Compute quantile-based bin boundaries for `PiecewiseLinearEncoding`/`PiecewiseLinearEmbeddings`. + +Note: `X` should have shape `(n_samples, n_features)`, which is the transpose of the +model input convention `(n_features, batch)`. Transpose your data before calling this. + +# Arguments +- `X::AbstractMatrix`: Training data of shape `(n_samples, n_features)`. +- `n_bins::Union{Int, Vector{Int}}`: Number of bins per feature (default `48`). + A single `Int` applies the same count to all features. A `Vector{Int}` of length + `n_features` specifies per-feature bin counts. + +# Returns +- `Vector{Vector{Float32}}`: Bin edges for each feature. Each vector has between 2 and + `n_bins + 1` elements (fewer if quantiles coincide for low-cardinality features). +""" +function compute_bins(X::AbstractMatrix; n_bins::Union{Int, Vector{Int}}=48) + n_samples, n_features = size(X) + + n_bins_vec = if n_bins isa Int + @assert n_bins > 1 "n_bins must be > 1, got $n_bins" + @assert n_bins < n_samples "n_bins must be < n_samples, got n_bins=$n_bins, n_samples=$n_samples" + fill(n_bins, n_features) + else + @assert length(n_bins) == n_features "Length of n_bins must match n_features ($n_features), got $(length(n_bins))" + for (j, nb) in enumerate(n_bins) + @assert nb > 1 "n_bins[$j] must be > 1, got $nb" + @assert nb < n_samples "n_bins[$j] must be < n_samples, got n_bins=$nb, n_samples=$n_samples" + end + n_bins + end + + bins = Vector{Vector{Float32}}(undef, n_features) + col_buf = Vector{eltype(X)}(undef, n_samples) + + for j in 1:n_features + copyto!(col_buf, view(X, :, j)) + sort!(col_buf) + quantile_probs = range(0f0, 1f0, length=n_bins_vec[j] + 1) + edges = Float32[quantile(col_buf, p; sorted=true) for p in quantile_probs] + unique!(edges) + @assert length(edges) >= 2 "Feature $j has fewer than 2 unique bin edges" + bins[j] = edges + end + return bins +end \ No newline at end of file diff --git a/src/models/embeddings/embeddings.jl b/src/models/embeddings/embeddings.jl new file mode 100644 index 0000000..ec2c807 --- /dev/null +++ b/src/models/embeddings/embeddings.jl @@ -0,0 +1,21 @@ +module Embeddings + +using Lux +using LuxCore +using Random +using NNlib +using LuxLib: batched_matmul +import Statistics: quantile + +export NLinear, LinearEmbeddings, LinearReLUEmbeddings +export Periodic, PeriodicEmbeddings +export PiecewiseLinearEncoding, PiecewiseLinearEmbeddings +export compute_bins + +include("compute_bins.jl") +include("nlinear.jl") +include("linear.jl") +include("periodic.jl") +include("piecewise_linear.jl") + +end \ No newline at end of file diff --git a/src/models/embeddings/linear.jl b/src/models/embeddings/linear.jl new file mode 100644 index 0000000..65af699 --- /dev/null +++ b/src/models/embeddings/linear.jl @@ -0,0 +1,42 @@ +using Lux +using Random +using NNlib + +""" + LinearEmbeddings(n_features, d_embedding; activation=relu) + +Embeds each continuous feature via a learned affine transformation followed by +an activation: `activation(w_j * x_j + b_j)`. +Produces a `(d_embedding, n_features, batch)` tensor. + +# Arguments +- `n_features::Int`: Number of input features. +- `d_embedding::Int`: Embedding dimension per feature. +- `activation`: Activation function applied element-wise (default `relu`). + E.g. `relu`, `tanh`, `identity`. +""" +struct LinearEmbeddings{F} <: Lux.AbstractLuxLayer + n_features::Int + d_embedding::Int + activation::F +end + +function LinearEmbeddings(n_features::Int, d_embedding::Int; activation=NNlib.relu) + return LinearEmbeddings(n_features, d_embedding, activation) +end + +function Lux.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) + limit = Float32(l.d_embedding)^(-0.5f0) + weight = reshape((rand(rng, Float32, l.d_embedding, l.n_features) .* 2f0 .* limit) .- limit, + l.d_embedding, l.n_features, 1) + bias = reshape((rand(rng, Float32, l.d_embedding, l.n_features) .* 2f0 .* limit) .- limit, + l.d_embedding, l.n_features, 1) + return (weight=weight, bias=bias) +end + +Lux.initialstates(::AbstractRNG, ::LinearEmbeddings) = (;) + +function (l::LinearEmbeddings)(x::AbstractMatrix, ps, st) + x_r = reshape(x, 1, size(x, 1), size(x, 2)) + return l.activation.(muladd.(ps.weight, x_r, ps.bias)), st +end \ No newline at end of file diff --git a/src/models/embeddings/nlinear.jl b/src/models/embeddings/nlinear.jl new file mode 100644 index 0000000..42f97b8 --- /dev/null +++ b/src/models/embeddings/nlinear.jl @@ -0,0 +1,50 @@ +using Lux +using Random +using NNlib + +""" + NLinear(n, in_features, out_features; bias=true) + +A batch of `n` independent linear layers applied in parallel via `batched_mul`. +Input shape `(in_features, n, batch)` → output shape `(out_features, n, batch)`. + +# Arguments +- `n::Int`: Number of independent linear layers (typically one per feature). +- `in_features::Int`: Input dimension for each linear layer. +- `out_features::Int`: Output dimension for each linear layer. +- `bias::Bool`: Whether to include a learnable bias (default `true`). +""" +struct NLinear <: Lux.AbstractLuxLayer + n::Int + in_features::Int + out_features::Int + use_bias::Bool +end + +function NLinear(n::Int, in_features::Int, out_features::Int; bias::Bool=true) + return NLinear(n, in_features, out_features, bias) +end + +function Lux.initialparameters(rng::AbstractRNG, l::NLinear) + limit = Float32(l.in_features)^(-0.5f0) + weight = (rand(rng, Float32, l.out_features, l.in_features, l.n) .* 2f0 .* limit) .- limit + + if l.use_bias + return (weight=weight, bias=zeros(Float32, l.out_features, 1, l.n)) + else + return (weight=weight,) + end +end + +Lux.initialstates(::AbstractRNG, ::NLinear) = (;) + +function (l::NLinear)(x::AbstractArray{T, 3}, ps, st) where T + x_perm = PermutedDimsArray(x, (1, 3, 2)) + out = batched_mul(ps.weight, x_perm) + + if l.use_bias + out = out .+ ps.bias + end + + return permutedims(out, (1, 3, 2)), st +end \ No newline at end of file diff --git a/src/models/embeddings/periodic.jl b/src/models/embeddings/periodic.jl new file mode 100644 index 0000000..73fbdaa --- /dev/null +++ b/src/models/embeddings/periodic.jl @@ -0,0 +1,95 @@ +using Lux +using Random +using NNlib + +""" + Periodic(n_features, n_frequencies, sigma) + +Maps each feature to `2 * n_frequencies` sinusoidal components: `[cos(2π w x), sin(2π w x)]`. +Output shape `(2 * n_frequencies, n_features, batch)`. + +# Arguments +- `n_features::Int`: Number of input features. +- `n_frequencies::Int`: Number of frequency components per feature. +- `sigma::Float32`: Std-dev for the frequency weight initialization (clamped to ±3σ). +""" +struct Periodic <: Lux.AbstractLuxLayer + n_features::Int + n_frequencies::Int + sigma::Float32 +end + +function Lux.initialparameters(rng::AbstractRNG, l::Periodic) + bound = l.sigma * 3f0 + w = clamp.(l.sigma .* randn(rng, Float32, l.n_frequencies, l.n_features), -bound, bound) + w = reshape(2f0 * Float32(π) .* w, l.n_frequencies, l.n_features, 1) + return (weight=w,) +end + +Lux.initialstates(::AbstractRNG, ::Periodic) = (;) + +function (l::Periodic)(x::AbstractMatrix, ps, st) + x_r = reshape(x, 1, size(x, 1), size(x, 2)) + z = ps.weight .* x_r + return vcat(cos.(z), sin.(z)), st +end + +""" + PeriodicEmbeddings(n_features, d_embedding=24; n_frequencies=48, + frequency_init_scale=0.01f0, activation=relu, lite=false) + +Periodic sinusoidal encoding followed by a learned linear projection. +Applies `Periodic` → `NLinear` (or `Dense` if `lite`) → activation. + +# Arguments +- `n_features::Int`: Number of input features. +- `d_embedding::Int`: Output embedding dimension per feature (default `24`). +- `n_frequencies::Int`: Sinusoidal frequency components per feature (default `48`). +- `frequency_init_scale::Float32`: σ for frequency weight init (default `0.01f0`). +- `activation`: Activation function applied after projection (default `relu`). E.g. `relu`, `tanh`, `identity`. +- `lite::Bool`: Use a single shared `Dense` instead of per-feature `NLinear` (default `false`). + Only valid when `activation` is not `identity`. +""" +struct PeriodicEmbeddings{P,L,F} <: Lux.AbstractLuxContainerLayer{(:periodic, :linear)} + periodic::P + linear::L + activation::F + lite::Bool +end + +function PeriodicEmbeddings( + n_features::Int, + d_embedding::Int=24; + n_frequencies::Int=48, + frequency_init_scale::Float32=0.01f0, + activation=NNlib.relu, + lite::Bool=false, +) + if lite && activation === identity + error("lite=true requires a non-identity activation function") + end + periodic = Periodic(n_features, n_frequencies, frequency_init_scale) + linear = if lite + Dense(2 * n_frequencies => d_embedding) + else + NLinear(n_features, 2 * n_frequencies, d_embedding) + end + return PeriodicEmbeddings(periodic, linear, activation, lite) +end + +function (m::PeriodicEmbeddings)(x::AbstractMatrix, ps, st) + h, st_p = m.periodic(x, ps.periodic, st.periodic) + + h_lin, st_l = if m.lite + d_in, n, b = size(h) + h_flat = reshape(h, d_in, n * b) + out_flat, st_sub = m.linear(h_flat, ps.linear, st.linear) + reshape(out_flat, size(out_flat, 1), n, b), st_sub + else + m.linear(h, ps.linear, st.linear) + end + + h_lin = m.activation.(h_lin) + + return h_lin, (periodic=st_p, linear=st_l) +end \ No newline at end of file diff --git a/src/models/embeddings/piecewise_linear.jl b/src/models/embeddings/piecewise_linear.jl new file mode 100644 index 0000000..b20cdaf --- /dev/null +++ b/src/models/embeddings/piecewise_linear.jl @@ -0,0 +1,137 @@ +using Lux +using Random +using NNlib + +""" + PiecewiseLinearEncoding(bins) + +Non-trainable piecewise-linear encoding using precomputed bin edges. +Output shape `(max_n_bins, n_features, batch)`. + +# Arguments +- `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). +""" +struct PiecewiseLinearEncoding <: Lux.AbstractLuxLayer + bins::Vector{Vector{Float32}} + n_features::Int + max_n_bins::Int +end + +function PiecewiseLinearEncoding(bins::Vector{<:AbstractVector}) + @assert length(bins) > 0 + n_features = length(bins) + n_bins_list = [length(b) - 1 for b in bins] + max_n_bins = maximum(n_bins_list) + bins_f32 = [Float32.(b) for b in bins] + return PiecewiseLinearEncoding(bins_f32, n_features, max_n_bins) +end + +Lux.initialparameters(::AbstractRNG, ::PiecewiseLinearEncoding) = (;) + +function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) + M, N = l.max_n_bins, l.n_features + + weight = zeros(Float32, M, N) + bias = zeros(Float32, M, N) + + for (i, bin_edges) in enumerate(l.bins) + bin_width = diff(bin_edges) + w = 1f0 ./ bin_width + b = -bin_edges[1:end-1] ./ bin_width + nb = length(bin_edges) - 1 + + # Place the last bin's weight/bias at the end row; + # remaining bins fill rows 1:nb-1. Unused rows stay zero + # and are clamped to [0, 1] harmlessly. + weight[end, i] = w[end] + bias[end, i] = b[end] + if nb > 1 + weight[1:nb-1, i] = w[1:end-1] + bias[1:nb-1, i] = b[1:end-1] + end + end + + # Pre-reshape to 3D to avoid per-call reshape in forward pass + return ( + weight = reshape(weight, M, N, 1), + bias = reshape(bias, M, N, 1), + ) +end + +function (l::PiecewiseLinearEncoding)(x::AbstractMatrix, ps, st) + # For each feature j and bin b: h[b,j,:] = clamp(w[b,j]*x[j,:] + bias[b,j], 0, 1) + x_r = reshape(x, 1, size(x, 1), size(x, 2)) + h = clamp.(muladd.(st.weight, x_r, st.bias), 0f0, 1f0) + return h, st +end + +""" + PiecewiseLinearEmbeddings(bins, d_embedding; activation=identity, version=:B) + +Learnable embeddings on top of `PiecewiseLinearEncoding`. +Version `:A`: PLE -> NLinear (with bias). +Version `:B`: PLE -> NLinear (zero-init, no bias) + LinearEmbeddings residual. +Output shape `(d_embedding, n_features, batch)`. + +# Arguments +- `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). +- `d_embedding::Int`: Embedding dimension per feature. +- `activation`: Activation function applied after projection (default `identity`). E.g. `relu`, `tanh`. +- `version::Symbol`: `:A` or `:B` (default `:B`). +""" +struct PiecewiseLinearEmbeddings{L0,I,L,F} <: Lux.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)} + linear0::L0 + encoding::I + linear::L + activation::F + version::Symbol +end + +function PiecewiseLinearEmbeddings( + bins::Vector{<:AbstractVector}, + d_embedding::Int; + activation=identity, + version::Symbol=:B, +) + @assert version in (:A, :B) + n_features = length(bins) + max_n_bins = maximum(length(b) - 1 for b in bins) + + encoding = PiecewiseLinearEncoding(bins) + # Residual path uses raw affine output (no activation) + linear0 = (version == :B) ? LinearEmbeddings(n_features, d_embedding; activation=identity) : nothing + linear = NLinear(n_features, max_n_bins, d_embedding; bias=(version == :A)) + + return PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version) +end + +function Lux.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbeddings) + ps_l0 = m.linear0 === nothing ? nothing : Lux.initialparameters(rng, m.linear0) + ps_enc = Lux.initialparameters(rng, m.encoding) + + if m.version == :B + n = m.linear.n + ps_lin = (weight=zeros(Float32, m.linear.out_features, m.linear.in_features, n),) + else + ps_lin = Lux.initialparameters(rng, m.linear) + end + + return (linear0=ps_l0, encoding=ps_enc, linear=ps_lin) +end + +function (m::PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st) + val_linear0 = nothing + st_l0 = st.linear0 + + if m.linear0 !== nothing + val_linear0, st_l0 = m.linear0(x, ps.linear0, st.linear0) + end + + h_enc, st_enc = m.encoding(x, ps.encoding, st.encoding) + h_proj, st_lin = m.linear(h_enc, ps.linear, st.linear) + + h_final = val_linear0 === nothing ? h_proj : (val_linear0 .+ h_proj) + h_final = m.activation.(h_final) + + return h_final, (linear0=st_l0, encoding=st_enc, linear=st_lin) +end \ No newline at end of file diff --git a/src/models/models.jl b/src/models/models.jl index 61033ac..354e607 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -1,27 +1,38 @@ module Models export NeuroTabModel, Architecture -export NeuroTreeConfig, MLPConfig, ResNetConfig +export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig +export Embeddings using ..Losses using Lux: Chain +using NNlib abstract type Architecture end +_broadcast_relu(x) = NNlib.relu.(x) + """ NeuroTabModel """ -struct NeuroTabModel{L<:LossType,C<:Chain} +struct NeuroTabModel{L<:LossType,C} _loss_type::Type{L} chain::C info::Dict{Symbol,Any} end # @functor NeuroTabModel (chain,) +include("embeddings/embeddings.jl") +using .Embeddings include("NeuroTree/neurotrees.jl") using .NeuroTrees + +include("TabM/TabM.jl") +using .TabM + # include("MLP/mlp.jl") # using .MLP + # include("ResNet/resnet.jl") # using .ResNet diff --git a/test/core.jl b/test/core.jl index 096f250..335b83c 100644 --- a/test/core.jl +++ b/test/core.jl @@ -1,5 +1,4 @@ @testset "Core - data iterators" begin - end @testset "Core - internals test" begin @@ -24,7 +23,7 @@ end nobs = 1_000 nfeats = 10 x = rand(Float32, nfeats, nobs) - feature_names = "var_" .* string.(1:nobs) + feature_names = "var_" .* string.(1:nfeats) outsize = 1 loss = NeuroTabModels.Losses.get_loss_fn(learner.loss) @@ -36,10 +35,10 @@ end ) m = NeuroTabModel(L, chain, info) - + end -@testset "Core - Regression" begin +@testset "Regression - NeuroTree" begin Random.seed!(123) X, y = rand(1000, 10), randn(1000) @@ -81,7 +80,48 @@ end end -@testset "Classification test" begin +@testset "Regression - TabM $at" for at in [:tabm, :tabm_mini, :tabm_packed] + + Random.seed!(123) + X = randn(Float32, 1000, 10) + y = X[:, 1] .+ 0.5f0 .* X[:, 2] .+ 0.1f0 .* randn(Float32, 1000) + df = DataFrame(X, :auto) + df[!, :y] = y + target_name = "y" + feature_names = setdiff(names(df), [target_name]) + + train_indices = 1:800 + dtrain = df[train_indices, :] + deval = df[801:end, :] + + arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=2, d_block=32, dropout=0.0, arch_type=at) + learner = NeuroTabRegressor(arch; loss=:mse, nrounds=20, early_stopping_rounds=2, lr=1e-2) + + m = NeuroTabModels.fit( + learner, + dtrain; + target_name, + feature_names + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + target_name, + feature_names, + deval, + ) + + p = m(deval) + @test size(p, 1) == nrow(deval) + @test !any(isnan, p) + mse_model = mean((p .- deval.y) .^ 2) + mse_baseline = mean((mean(dtrain.y) .- deval.y) .^ 2) + @test mse_model < mse_baseline + +end + +@testset "Classification - NeuroTree" begin Random.seed!(123) X, y = @load_crabs @@ -113,7 +153,6 @@ end target_name, feature_names, ) - # Predictions depend on the number of samples in the dataset ptrain = [argmax(x) for x in eachrow(m(dtrain))] peval = [argmax(x) for x in eachrow(m(deval))] @@ -121,3 +160,41 @@ end @test mean(peval .== levelcode.(deval.class)) > 0.95 end + +@testset "Classification - TabM $at" for at in [:tabm, :tabm_mini, :tabm_packed] + + Random.seed!(123) + X, y = @load_crabs + df = DataFrame(X) + df[!, :class] = y + target_name = "class" + feature_names = setdiff(names(df), [target_name]) + + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + arch = NeuroTabModels.TabMConfig(; k=2, n_blocks=2, d_block=16, dropout=0.0, arch_type=at) + learner = NeuroTabClassifier(arch; + nrounds=500, + batchsize=32, + early_stopping_rounds=50, + lr=1e-2, + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + deval, + target_name, + feature_names, + ) + + ptrain = [argmax(x) for x in eachrow(m(dtrain))] + peval = [argmax(x) for x in eachrow(m(deval))] + @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 + @test mean(peval .== levelcode.(deval.class)) > 0.95 + +end \ No newline at end of file diff --git a/test/embedding.jl b/test/embedding.jl new file mode 100644 index 0000000..bfeb261 --- /dev/null +++ b/test/embedding.jl @@ -0,0 +1,42 @@ +@testset "Numerical Embeddings" begin + + Random.seed!(123) + n = 1000 + X = randn(Float32, n, 10) + y = X[:, 1] .+ 0.5f0 .* X[:, 2] .+ 0.1f0 .* randn(Float32, n) + + df = DataFrame(X, :auto) + df[!, :y] = y + target_name = "y" + feature_names = setdiff(names(df), [target_name]) + + train_indices = 1:800 + dtrain = df[train_indices, :] + deval = df[801:end, :] + + mse_baseline = mean((mean(dtrain.y) .- deval.y) .^ 2) + + # (name, base config kwargs) + architectures = [ + ("TabM", Dict( + :k => 4, :n_blocks => 2, :d_block => 32, :dropout => 0.0)), + ] + + @testset "$arch_name - $et" for (arch_name, base_kw) in architectures, + et in [:periodic, :linear, :piecewise] + + extra = et == :piecewise ? Dict(:n_bins => 16) : Dict() + arch = NeuroTabModels.TabMConfig(; + base_kw..., use_embeddings=true, d_embedding=8, embedding_type=et, extra...) + + learner = NeuroTabRegressor(arch; loss=:mse, nrounds=20, lr=1e-2) + m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names) + + p = m(deval) + @test size(p, 1) == nrow(deval) + @test !any(isnan, p) + @test mean((p .- deval.y) .^ 2) < mse_baseline + + end + +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index b93c8bd..8959b2b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,4 +10,5 @@ using MLJBase using MLJTestInterface include("core.jl") +include("embedding.jl") include("MLJ.jl") From 9456f70d7080830c163e7c86d3f83d92fa157d28 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Thu, 19 Mar 2026 01:31:43 -0400 Subject: [PATCH 083/120] docs --- docs/make.jl | 7 +++++-- docs/src/{ => models}/models.md | 0 docs/src/models/neurotrees.md | 5 +++++ docs/src/models/tabM.md | 5 +++++ src/Fit/fit.jl | 5 +++-- src/NeuroTabModels.jl | 1 - src/losses.jl | 10 +++++----- src/models/TabM/TabM.jl | 20 ++++++++++---------- src/models/models.jl | 9 ++++++++- 9 files changed, 41 insertions(+), 21 deletions(-) rename docs/src/{ => models}/models.md (100%) create mode 100644 docs/src/models/neurotrees.md create mode 100644 docs/src/models/tabM.md diff --git a/docs/make.jl b/docs/make.jl index 868e8fa..8670f59 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -6,8 +6,11 @@ using NeuroTabModels pages = [ "Quick start" => "quick-start.md", "Design" => "design.md", - "Models" => "models.md", - "API" => "API.md", + "Models" => [ + "Interface" => "models/models.md", + "NeuroTrees" => "models/neurotrees.md", + "TabM" => "models/tabM.md", + ], "Tutorials" => [ "Regression - Boston" => "tutorials/regression-boston.md", "Logistic - Titanic" => "tutorials/logistic-titanic.md", diff --git a/docs/src/models.md b/docs/src/models/models.md similarity index 100% rename from docs/src/models.md rename to docs/src/models/models.md diff --git a/docs/src/models/neurotrees.md b/docs/src/models/neurotrees.md new file mode 100644 index 0000000..1787c2a --- /dev/null +++ b/docs/src/models/neurotrees.md @@ -0,0 +1,5 @@ +# NeuroTrees + +```@docs +NeuroTreeConfig +``` diff --git a/docs/src/models/tabM.md b/docs/src/models/tabM.md new file mode 100644 index 0000000..e3300b6 --- /dev/null +++ b/docs/src/models/tabM.md @@ -0,0 +1,5 @@ +# TabM + +```@docs +TabMConfig +``` diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index da2ece6..b254c50 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -83,7 +83,7 @@ function init( end """ - function fit( + fit( config::LearnerTypes, dtrain; feature_names, @@ -96,7 +96,9 @@ end early_stopping_rounds=9999, verbosity=1, ) + Training function of NeuroTabModels' internal API. + # Arguments - `config::LearnerTypes`: The configuration object defining the model architecture, loss, and training hyperparameters. @@ -116,7 +118,6 @@ Training function of NeuroTabModels' internal API. - `device=:cpu`: Symbol. Hardware device to use for training (`:cpu` or `:gpu`). - `gpuID=0`: Integer. Specifies which GPU to use if multiple are available. """ - function fit( config::LearnerTypes, dtrain; diff --git a/src/NeuroTabModels.jl b/src/NeuroTabModels.jl index a13c884..a939477 100644 --- a/src/NeuroTabModels.jl +++ b/src/NeuroTabModels.jl @@ -3,7 +3,6 @@ module NeuroTabModels using Random export NeuroTabRegressor, NeuroTabClassifier, NeuroTabModel -export TabMConfig, compute_bins include("data.jl") include("losses.jl") diff --git a/src/losses.jl b/src/losses.jl index 2648d47..b3a0ef6 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -70,16 +70,16 @@ end function gaussian_mle(model, ps, st, data::Tuple{Any,Any}) pred, st_ = _forward(model, ps, st, data[1]) - return _reduce(_gaussian_mle_core(pred[1:1,:,:], pred[2:2,:,:], _reshape_3d(data[2]))), st_, NamedTuple() + return _reduce(_gaussian_mle_core(pred[1:1, :, :], pred[2:2, :, :], _reshape_3d(data[2]))), st_, NamedTuple() end function gaussian_mle(model, ps, st, data::Tuple{Any,Any,Any}) pred, st_ = _forward(model, ps, st, data[1]) - return _reduce(_gaussian_mle_core(pred[1:1,:,:], pred[2:2,:,:], _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() + return _reduce(_gaussian_mle_core(pred[1:1, :, :], pred[2:2, :, :], _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() end function gaussian_mle(model, ps, st, data::Tuple{Any,Any,Any,Any}) pred, st_ = _forward(model, ps, st, data[1]) pred = pred .+ _reshape_3d(data[4]) - return _reduce(_gaussian_mle_core(pred[1:1,:,:], pred[2:2,:,:], _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() + return _reduce(_gaussian_mle_core(pred[1:1, :, :], pred[2:2, :, :], _reshape_3d(data[2])), _reshape_3d(data[3])), st_, NamedTuple() end get_loss_fn(::Type{<:MSE}) = mse_loss @@ -89,12 +89,12 @@ get_loss_fn(::Type{<:MLogLoss}) = mlogloss get_loss_fn(::Type{<:GaussianMLE}) = gaussian_mle get_loss_fn(::Type{<:Tweedie}) = tweedie -const _loss_type_dict = Dict( +const loss_type_dict = Dict( :mse => MSE, :mae => MAE, :logloss => LogLoss, :mlogloss => MLogLoss, :gaussian_mle => GaussianMLE, :tweedie => Tweedie, ) -get_loss_type(loss::Symbol) = _loss_type_dict[loss] +get_loss_type(loss::Symbol) = loss_type_dict[loss] get_loss_fn(s::Symbol) = get_loss_fn(get_loss_type(s)) end \ No newline at end of file diff --git a/src/models/TabM/TabM.jl b/src/models/TabM/TabM.jl index df4b386..6a38131 100644 --- a/src/models/TabM/TabM.jl +++ b/src/models/TabM/TabM.jl @@ -15,18 +15,18 @@ import ..Models: _broadcast_relu include("layers.jl") function _batch_ensemble_backbone(; - d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, - k::Int, scaling_init::Symbol, d_features::Vector{Int}) + d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, + k::Int, scaling_init::Symbol, d_features::Vector{Int}) layers = Any[] for i in 1:n_blocks d_in_i = (i == 1) ? d_in : d_block if i == 1 push!(layers, LinearBatchEnsemble(d_in_i, d_block; - k, scaling_init = (scaling_init, :ones), - first_scaling_init_chunks = d_features)) + k, scaling_init=(scaling_init, :ones), + first_scaling_init_chunks=d_features)) else push!(layers, LinearBatchEnsemble(d_in_i, d_block; - k, scaling_init = :ones)) + k, scaling_init=:ones)) end push!(layers, WrappedFunction(_broadcast_relu)) dropout > 0 && push!(layers, Dropout(dropout)) @@ -35,10 +35,10 @@ function _batch_ensemble_backbone(; end function _mini_ensemble_backbone(; - d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, - k::Int, scaling_init::Symbol, d_features::Vector{Int}) + d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, + k::Int, scaling_init::Symbol, d_features::Vector{Int}) layers = Any[ScaleEnsemble(k, d_in; - init = scaling_init, init_chunks = d_features, bias = false)] + init=scaling_init, init_chunks=d_features, bias=false)] for i in 1:n_blocks d_in_i = (i == 1) ? d_in : d_block push!(layers, Dense(d_in_i => d_block, relu)) @@ -48,7 +48,7 @@ function _mini_ensemble_backbone(; end function _packed_ensemble_backbone(; - d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, k::Int) + d_in::Int, n_blocks::Int, d_block::Int, dropout::Float64, k::Int) layers = Any[] for i in 1:n_blocks d_in_i = (i == 1) ? d_in : d_block @@ -104,7 +104,7 @@ struct TabMConfig <: Architecture use_embeddings::Bool d_embedding::Int embedding_type::Symbol - n_bins::Union{Int, Vector{Int}} + n_bins::Union{Int,Vector{Int}} end function TabMConfig(; kwargs...) diff --git a/src/models/models.jl b/src/models/models.jl index 354e607..064daaa 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -14,9 +14,16 @@ _broadcast_relu(x) = NNlib.relu.(x) """ NeuroTabModel + +A trained neural network model for tabular data. It is the object returned by [`fit`](@ref) and wraps a [Lux.jl](https://lux.csail.mit.edu) `chain` built from one of the supported [`Architecture`](@ref) configurations: +## Fields + +- `loss_type`: the loss function type used during training (e.g. `MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`) +- `chain`: the underlying `Lux.Chain` neural network +- `info`: a `Dict{Symbol,Any}` storing metadata such as `:feature_names`, `:target_levels`, and `:device` """ struct NeuroTabModel{L<:LossType,C} - _loss_type::Type{L} + loss_type::Type{L} chain::C info::Dict{Symbol,Any} end From 439f5cd6371fbe223c0e43149da9da2e489c4859 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Thu, 19 Mar 2026 01:55:11 -0400 Subject: [PATCH 084/120] docs --- benchmarks/titanic-mlogloss.jl | 15 ++++++++++++++- docs/src/models/models.md | 3 +++ src/models/models.jl | 6 +++++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/benchmarks/titanic-mlogloss.jl b/benchmarks/titanic-mlogloss.jl index 74070f8..e3cae01 100644 --- a/benchmarks/titanic-mlogloss.jl +++ b/benchmarks/titanic-mlogloss.jl @@ -5,7 +5,6 @@ using Statistics: mean using StatsBase: median using CategoricalArrays using Random -using CUDA using CategoricalArrays Random.seed!(123) @@ -45,11 +44,25 @@ arch = NeuroTabModels.NeuroTreeConfig(; hidden_size=1, ) +# arch = NeuroTabModels.TabMConfig(; +# arch_type=:tabm_mini, +# k=16, +# d_block=32, +# n_blocks=1, +# dropout=0.1, +# use_embeddings=true, +# embedding_type=:piecewise, +# n_bins=32, +# d_embedding=32, +# scaling_init=:normal, +# ) + learner = NeuroTabClassifier( arch; nrounds=100, early_stopping_rounds=2, lr=1e-2, + device=:cpu, ) m = NeuroTabModels.fit( diff --git a/docs/src/models/models.md b/docs/src/models/models.md index 6fd6ec7..3db7340 100644 --- a/docs/src/models/models.md +++ b/docs/src/models/models.md @@ -1,16 +1,19 @@ # Models ## NeuroTabRegressor + ```@docs NeuroTabRegressor ``` ## NeuroTabClassifier + ```@docs NeuroTabClassifier ``` ## NeuroTabModel + ```@docs NeuroTabModel ``` diff --git a/src/models/models.jl b/src/models/models.jl index 064daaa..2cff884 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -15,7 +15,11 @@ _broadcast_relu(x) = NNlib.relu.(x) """ NeuroTabModel -A trained neural network model for tabular data. It is the object returned by [`fit`](@ref) and wraps a [Lux.jl](https://lux.csail.mit.edu) `chain` built from one of the supported [`Architecture`](@ref) configurations: +A trained neural network model for tabular data. It is the object returned by [`NeuroTabModels.fit`](@ref) and wraps a [Lux.jl](https://lux.csail.mit.edu) `chain` built from one of the supported `Architecture` configurations: + + - NeuroTreeConfig + - TabMConfig + ## Fields - `loss_type`: the loss function type used during training (e.g. `MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`) From c0d766c112f9a28d8fac966d616b03914782f6fc Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Thu, 19 Mar 2026 02:17:00 -0400 Subject: [PATCH 085/120] docs --- src/models/models.jl | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/models/models.jl b/src/models/models.jl index 2cff884..84ccf56 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -15,16 +15,13 @@ _broadcast_relu(x) = NNlib.relu.(x) """ NeuroTabModel -A trained neural network model for tabular data. It is the object returned by [`NeuroTabModels.fit`](@ref) and wraps a [Lux.jl](https://lux.csail.mit.edu) `chain` built from one of the supported `Architecture` configurations: - - - NeuroTreeConfig - - TabMConfig +The object containing the model and associated metadata. ## Fields - `loss_type`: the loss function type used during training (e.g. `MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`) - `chain`: the underlying `Lux.Chain` neural network -- `info`: a `Dict{Symbol,Any}` storing metadata such as `:feature_names`, `:target_levels`, and `:device` +- `info`: a `Dict{Symbol,Any}` storing metadata such as `:feature_names`, `:target_levels`, `:device`, `logger`, as well as fitted parameters (`ps`) and state (`st`). """ struct NeuroTabModel{L<:LossType,C} loss_type::Type{L} From c9a9747c25e705aadaac5aa7c0202a3e0e54f90b Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Wed, 25 Mar 2026 10:24:15 +0530 Subject: [PATCH 086/120] Add EmbeddingConfig for architecture-independent embeddings (#33) * - Add EmbeddingConfig as architecture-agnostic embedding configuration - Move embedding construction from TabMConfig to fit.jl (_build_chain) - Embeddings now specified via embedding_config kwarg on learner constructors - Support Dict -> EmbeddingConfig conversion for ergonomic API - Add activation function parameter to LinearEmbeddings, PeriodicEmbeddings, PiecewiseLinearEmbeddings (replaces hardcoded relu/bool) - Fix PiecewiseLinearEncoding to use clamp(0,1) per paper specification - Unify _ensure_3d/_bcast into _reshape_3d in losses.jl - Rename _activation to _inverse_link in infer.jl - Remove LinearReLUEmbeddings (redundant with LinearEmbeddings(activation=relu)) - Add embedding regression + classification tests Signed-off-by: AdityaPandeyCN * tune test params Signed-off-by: AdityaPandeyCN * Simplify chain construction and restore scaling_init override Signed-off-by: AdityaPandeyCN * resolve scaling_init issue Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN --- src/Fit/callback.jl | 2 +- src/Fit/fit.jl | 22 ++++-- src/learners.jl | 36 ++++++++-- src/models/NeuroTree/neurotrees.jl | 2 +- src/models/TabM/TabM.jl | 100 ++++++++++------------------ src/models/embeddings/config.jl | 68 +++++++++++++++++++ src/models/embeddings/embeddings.jl | 6 +- test/embedding.jl | 61 ++++++++++++++--- 8 files changed, 205 insertions(+), 92 deletions(-) create mode 100644 src/models/embeddings/config.jl diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 8ac5f76..11b1124 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -4,7 +4,7 @@ using DataFrames using Statistics: mean, median using ..Learners: LearnerTypes -using ..Infer: reduce_pred +using ...Infer: reduce_pred using ..Data: get_df_loader_train using ..Metrics diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index b254c50..b86c54d 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -7,7 +7,7 @@ using ..Learners using ..Models using ..Losses using ..Metrics -using ..Infer +using ..Infer: reduce_pred import Random: Xoshiro import MLJModelInterface: fit @@ -66,12 +66,22 @@ function init( :device => config.device ) - if hasproperty(config.arch, :use_embeddings) && config.arch.use_embeddings && config.arch.embedding_type == :piecewise - X_train = Matrix{Float32}(df[:, feature_names]) - chain = config.arch(; nfeats, outsize, X_train) - else + # Build chain: optional embeddings + architecture backbone + embed_config = config.embedding_config + if isnothing(embed_config) chain = config.arch(; nfeats, outsize) + else + if embed_config.embedding_type == :piecewise + X_train = Matrix{Float32}(df[:, feature_names]) + else + X_train = nothing + end + embed_chain = embed_config(; nfeats, X_train) + d_in = nfeats * embed_config.d_embedding + d_features = fill(embed_config.d_embedding, nfeats) + chain = Chain(embed_chain, config.arch(; nfeats=d_in, outsize, d_features, scaling_init_override=:normal)) end + m = NeuroTabModel(L, chain, info) rng = Xoshiro(config.seed) @@ -111,7 +121,7 @@ Training function of NeuroTabModels' internal API. - `weight_name=nothing`: Optional. A `Symbol` or `String` indicating the sample weights column. - `offset_name=nothing`: Optional. A `Symbol` or `String` indicating the offset column. - `deval=nothing`: Optional. Evaluation data (`<:AbstractDataFrame`) for tracking metrics and early stopping. -- `metric=nothing`: Optional. The evaluation metric to track (e.g., `:mse`, `:logloss`). +- `metric=nothing`: Optional. The evaluation metric to track (e.g., `:mse`, `:logloss`). - `print_every_n=9999`: Integer. Logs training progress to the console every `N` epochs. - `early_stopping_rounds=9999`: Integer. Stops training if the evaluation metric does not improve for this many rounds. - `verbosity=1`: Integer. Controls the logging level (`0` for silent, `>0` for info). diff --git a/src/learners.jl b/src/learners.jl index 4a38129..7bc72d5 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -5,12 +5,14 @@ import MLJModelInterface: fit, update, predict, schema import Random using ..Models +using ..Models: EmbeddingConfig export NeuroTabRegressor, NeuroTabClassifier, LearnerTypes mutable struct NeuroTabRegressor <: MMI.Deterministic loss::Symbol metric::Symbol arch::Architecture + embedding_config::Union{Nothing,EmbeddingConfig} nrounds::Int early_stopping_rounds::Int lr::Float32 @@ -36,12 +38,14 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl]( - `:mlogloss` - `:gaussian_mle` - `nrounds=100`: Max number of rounds (epochs). -- `lr=1.0f-2`: Learning rate. Must be > 0. A lower `eta` results in slower learning, typically requiring a higher `nrounds`. +- `lr=1.0f-2`: Learning rate. Must be > 0. A lower `eta` results in slower learning, typically requiring a higher `nrounds`. - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. - `device=:cpu`: Device on which to perform the computation, either `:cpu` or `:gpu` -- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. + E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`. # Internal API @@ -141,7 +145,8 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) :batchsize => 2048, :seed => 123, :device => :cpu, - :gpuID => 0 + :gpuID => 0, + :embedding_config => nothing, ) args_ignored = setdiff(keys(kwargs), keys(args)) @@ -174,10 +179,17 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) device = Symbol(args[:device]) + # Build EmbeddingConfig from Dict if needed + embed = args[:embedding_config] + if embed isa AbstractDict + embed = EmbeddingConfig(; embed...) + end + config = NeuroTabRegressor( loss, metric, arch, + embed, args[:nrounds], args[:early_stopping_rounds], Float32(args[:lr]), @@ -202,6 +214,7 @@ mutable struct NeuroTabClassifier <: MMI.Probabilistic loss::Symbol metric::Symbol arch::Architecture + embedding_config::Union{Nothing,EmbeddingConfig} nrounds::Int early_stopping_rounds::Int lr::Float32 @@ -221,12 +234,13 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl] # Hyper-parameters - `nrounds=100`: Max number of rounds (epochs). -- `lr=1.0f-2`: Learning rate. Must be > 0. A lower `eta` results in slower learning, typically requiring a higher `nrounds`. +- `lr=1.0f-2`: Learning rate. Must be > 0. A lower `eta` results in slower learning, typically requiring a higher `nrounds`. - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. - `device=:cpu`: Device on which to perform the computation, either `:cpu` or `:gpu` -- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. # Internal API @@ -293,7 +307,7 @@ The fields of `report(mach)` are: ## Internal API ```julia -using NeuroTabModels, DataFrames, CategoricalArrays, Random +using NeuroTabModels, DataFrames, CategoricalArrays, Random config = NeuroTabClassifier(depth=5, nrounds=10) nobs, nfeats = 1_000, 5 dtrain = DataFrame(randn(nobs, nfeats), :auto) @@ -325,7 +339,8 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) :batchsize => 2048, :seed => 123, :device => :cpu, - :gpuID => 0 + :gpuID => 0, + :embedding_config => nothing, ) args_ignored = setdiff(keys(kwargs), keys(args)) @@ -345,10 +360,17 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) device = Symbol(args[:device]) + # Build EmbeddingConfig from Dict if needed + embed = args[:embedding_config] + if embed isa AbstractDict + embed = EmbeddingConfig(; embed...) + end + config = NeuroTabClassifier( :mlogloss, :mlogloss, arch, + embed, args[:nrounds], args[:early_stopping_rounds], Float32(args[:lr]), diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 42860bb..d4b9a35 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -97,7 +97,7 @@ function _tree_kwargs(config::NeuroTreeConfig) ) end -function (config::NeuroTreeConfig)(; nfeats, outsize) +function (config::NeuroTreeConfig)(; nfeats, outsize, kwargs...) kwargs = _tree_kwargs(config) if config.MLE_tree_split diff --git a/src/models/TabM/TabM.jl b/src/models/TabM/TabM.jl index 6a38131..2e1abc0 100644 --- a/src/models/TabM/TabM.jl +++ b/src/models/TabM/TabM.jl @@ -8,9 +8,7 @@ using LuxCore using NNlib import ..Losses: get_loss_type, GaussianMLE -import ..Models: Architecture -import ..Embeddings: PeriodicEmbeddings, LinearEmbeddings, PiecewiseLinearEmbeddings, compute_bins -import ..Models: _broadcast_relu +import ..Models: Architecture, _broadcast_relu include("layers.jl") @@ -62,36 +60,22 @@ end """ TabMConfig(; kwargs...) -Configuration for TabM ensemble architectures (Gorishniy et al., ICLR 2025). +Configuration for TabM ensemble architectures. -This implements the TabM♠ variant (shared training batches), where all k ensemble -members see the same training batch. - -The chain output is 3D: `(outsize, k, batch)`. Ensemble averaging is handled by -`Losses.reduce_pred` during training and `_reduce` in `infer.jl` at inference. +The chain output is 3D: `(outsize, k, batch)`. Ensemble averaging is handled +during training and at inference. # Arguments - `k::Int`: Number of ensemble members (default `32`). -- `n_blocks::Int`: Number of MLP blocks (default `2` with embeddings, `3` without). +- `n_blocks::Int`: Number of MLP blocks (default `3`). - `d_block::Int`: Hidden dimension per block (default `512`). - `dropout::Float64`: Dropout rate (default `0.1`). - `arch_type::Symbol`: `:tabm`, `:tabm_mini`, or `:tabm_packed` (default `:tabm`). -- `scaling_init::Symbol`: Init for ensemble scaling vectors — `:random_signs`, `:normal`, or `:ones` (default `:normal` with embeddings, `:random_signs` without). +- `scaling_init::Union{Nothing,Symbol}`: Init for ensemble scaling — `:random_signs`, `:normal`, + `:ones`, or `nothing` for auto-detection (default `nothing`). + When `nothing`, defaults to `:random_signs` without embeddings and `:normal` with embeddings. + Explicit values always take priority. - `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). -- `use_embeddings::Bool`: Apply feature embeddings before the backbone (default `false`). -- `d_embedding::Int`: Embedding dimension per feature (default `24`). -- `embedding_type::Symbol`: `:periodic`, `:linear`, or `:piecewise` (default `:periodic`). -- `n_bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). - A single `Int` applies the same count to all features. A `Vector{Int}` specifies - per-feature bin counts. - -# Callable - config(; nfeats, outsize, X_train=nothing) → Lux.Chain - -- `nfeats::Int`: Number of input features. -- `outsize::Int`: Output dimension. -- `X_train::Union{Nothing, AbstractMatrix}`: Training data of shape `(n_samples, n_features)`. - Required when `embedding_type=:piecewise` to compute bin edges via `compute_bins`. """ struct TabMConfig <: Architecture k::Int @@ -99,29 +83,19 @@ struct TabMConfig <: Architecture d_block::Int dropout::Float64 arch_type::Symbol - scaling_init::Symbol + scaling_init::Union{Nothing,Symbol} MLE_tree_split::Bool - use_embeddings::Bool - d_embedding::Int - embedding_type::Symbol - n_bins::Union{Int,Vector{Int}} end function TabMConfig(; kwargs...) - has_embeddings = get(kwargs, :use_embeddings, false) - args = Dict{Symbol,Any}( :k => 32, - :n_blocks => has_embeddings ? 2 : 3, + :n_blocks => 3, :d_block => 512, :dropout => 0.1, :arch_type => :tabm, - :scaling_init => has_embeddings ? :normal : :random_signs, + :scaling_init => nothing, :MLE_tree_split => false, - :use_embeddings => false, - :d_embedding => 24, - :embedding_type => :periodic, - :n_bins => 48, ) args_ignored = setdiff(keys(kwargs), keys(args)) @@ -136,54 +110,48 @@ function TabMConfig(; kwargs...) args[arg] = kwargs[arg] end + si = args[:scaling_init] + scaling_init = isnothing(si) ? nothing : Symbol(si) + return TabMConfig( args[:k], args[:n_blocks], args[:d_block], args[:dropout], - Symbol(args[:arch_type]), Symbol(args[:scaling_init]), - args[:MLE_tree_split], args[:use_embeddings], - args[:d_embedding], Symbol(args[:embedding_type]), args[:n_bins], + Symbol(args[:arch_type]), scaling_init, + args[:MLE_tree_split], ) end -function (config::TabMConfig)(; nfeats, outsize, X_train=nothing) +function (config::TabMConfig)(; nfeats, outsize, d_features=nothing, scaling_init_override=nothing) @assert config.k > 0 "k must be > 0, got $(config.k)" @assert nfeats > 0 "nfeats must be > 0, got $nfeats" @assert outsize > 0 "outsize must be > 0, got $outsize" k = config.k d_block = config.d_block + d_in = nfeats - if config.use_embeddings - emb_layer = if config.embedding_type == :periodic - PeriodicEmbeddings(nfeats, config.d_embedding) - elseif config.embedding_type == :linear - LinearEmbeddings(nfeats, config.d_embedding) - elseif config.embedding_type == :piecewise - @assert X_train !== nothing "Piecewise embeddings require `X_train` to compute bin edges." - bins = compute_bins(X_train; n_bins=config.n_bins) - @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))" - PiecewiseLinearEmbeddings(bins, config.d_embedding) - else - error("Unsupported embedding type: $(config.embedding_type)") - end - feature_layers = [emb_layer, FlattenLayer()] - d_in = nfeats * config.d_embedding - d_features = fill(config.d_embedding, nfeats) - effective_scaling_init = :normal - else - feature_layers = [] - d_in = nfeats + if isnothing(d_features) d_features = ones(Int, nfeats) - effective_scaling_init = config.scaling_init + end + + # 1. User explicitly set it - use it + # 2. User left it as nothing, caller suggests override - use override + # 3. Neither - default to :random_signs + scaling_init = if !isnothing(config.scaling_init) + config.scaling_init + elseif !isnothing(scaling_init_override) + scaling_init_override + else + :random_signs end bb = if config.arch_type == :tabm _batch_ensemble_backbone(; d_in, n_blocks=config.n_blocks, d_block, dropout=config.dropout, k, - scaling_init=effective_scaling_init, d_features) + scaling_init, d_features) elseif config.arch_type == :tabm_mini _mini_ensemble_backbone(; d_in, n_blocks=config.n_blocks, d_block, dropout=config.dropout, k, - scaling_init=effective_scaling_init, d_features) + scaling_init, d_features) elseif config.arch_type == :tabm_packed _packed_ensemble_backbone(; d_in, n_blocks=config.n_blocks, d_block, dropout=config.dropout, k) @@ -200,7 +168,7 @@ function (config::TabMConfig)(; nfeats, outsize, X_train=nothing) LinearEnsemble(d_block, outsize, k) end - return Chain(feature_layers..., EnsembleView(k), bb..., head) + return Chain(EnsembleView(k), bb..., head) end end \ No newline at end of file diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl new file mode 100644 index 0000000..a5b8bb3 --- /dev/null +++ b/src/models/embeddings/config.jl @@ -0,0 +1,68 @@ +using NNlib: relu + +""" + EmbeddingConfig(; kwargs...) + +Architecture-agnostic configuration for numerical feature embeddings. +Constructed by the user and passed as `embedding_config` to `NeuroTabRegressor`/`NeuroTabClassifier`. + +# Arguments +- `embedding_type::Symbol`: `:periodic`, `:linear`, or `:piecewise` (default `:periodic`). +- `d_embedding::Int`: Embedding dimension per feature (default `24`). +- `activation`: Activation function after projection (default `relu` for periodic/linear, + `identity` for piecewise). +- `n_bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). +- `n_frequencies::Int`: Frequency components for periodic embeddings (default `48`). +- `frequency_init_scale::Float32`: σ for periodic frequency init (default `0.01f0`). + +# Callable + config(; nfeats, X_train=nothing)-Lux.Chain + +Returns a `Chain(embedding_layer, FlattenLayer())` ready to prepend to any backbone. +""" +struct EmbeddingConfig{F} + embedding_type::Symbol + d_embedding::Int + activation::F + n_bins::Union{Int, Vector{Int}} + n_frequencies::Int + frequency_init_scale::Float32 +end + +function EmbeddingConfig(; + embedding_type::Symbol=:periodic, + d_embedding::Int=24, + activation=nothing, + n_bins::Union{Int, Vector{Int}}=48, + n_frequencies::Int=48, + frequency_init_scale::Float32=0.01f0, +) + # Default activation depends on embedding type + if isnothing(activation) + activation = embedding_type == :piecewise ? identity : relu + end + + return EmbeddingConfig( + embedding_type, d_embedding, activation, + n_bins, n_frequencies, frequency_init_scale, + ) +end + +function (config::EmbeddingConfig)(; nfeats::Int, X_train=nothing) + emb = if config.embedding_type == :periodic + PeriodicEmbeddings(nfeats, config.d_embedding; + n_frequencies=config.n_frequencies, + frequency_init_scale=config.frequency_init_scale, + activation=config.activation) + elseif config.embedding_type == :linear + LinearEmbeddings(nfeats, config.d_embedding; activation=config.activation) + elseif config.embedding_type == :piecewise + @assert X_train !== nothing "Piecewise embeddings require `X_train` to compute bin edges." + bins = compute_bins(X_train; n_bins=config.n_bins) + @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))" + PiecewiseLinearEmbeddings(bins, config.d_embedding; activation=config.activation) + else + error("Unsupported embedding type: $(config.embedding_type)") + end + return Chain(emb, FlattenLayer()) +end \ No newline at end of file diff --git a/src/models/embeddings/embeddings.jl b/src/models/embeddings/embeddings.jl index ec2c807..b24f672 100644 --- a/src/models/embeddings/embeddings.jl +++ b/src/models/embeddings/embeddings.jl @@ -1,21 +1,23 @@ module Embeddings using Lux +using Lux: Chain, FlattenLayer using LuxCore using Random using NNlib using LuxLib: batched_matmul import Statistics: quantile -export NLinear, LinearEmbeddings, LinearReLUEmbeddings +export NLinear, LinearEmbeddings export Periodic, PeriodicEmbeddings export PiecewiseLinearEncoding, PiecewiseLinearEmbeddings -export compute_bins +export compute_bins, EmbeddingConfig include("compute_bins.jl") include("nlinear.jl") include("linear.jl") include("periodic.jl") include("piecewise_linear.jl") +include("config.jl") end \ No newline at end of file diff --git a/test/embedding.jl b/test/embedding.jl index bfeb261..06d494f 100644 --- a/test/embedding.jl +++ b/test/embedding.jl @@ -1,4 +1,4 @@ -@testset "Numerical Embeddings" begin +@testset "Embeddings - Regression" begin Random.seed!(123) n = 1000 @@ -16,20 +16,23 @@ mse_baseline = mean((mean(dtrain.y) .- deval.y) .^ 2) - # (name, base config kwargs) architectures = [ - ("TabM", Dict( - :k => 4, :n_blocks => 2, :d_block => 32, :dropout => 0.0)), + ("TabM", NeuroTabModels.TabMConfig(; k=4, n_blocks=2, d_block=32, dropout=0.0)), + ("NeuroTree", NeuroTabModels.NeuroTreeConfig(; depth=3)), ] - @testset "$arch_name - $et" for (arch_name, base_kw) in architectures, + @testset "$arch_name - $et" for (arch_name, arch) in architectures, et in [:periodic, :linear, :piecewise] - extra = et == :piecewise ? Dict(:n_bins => 16) : Dict() - arch = NeuroTabModels.TabMConfig(; - base_kw..., use_embeddings=true, d_embedding=8, embedding_type=et, extra...) + embed_kw = Dict(:embedding_type => et, :d_embedding => 8) + if et == :piecewise + embed_kw[:n_bins] = 16 + end + + learner = NeuroTabRegressor(arch; + loss=:mse, nrounds=20, lr=1e-2, + embedding_config=embed_kw) - learner = NeuroTabRegressor(arch; loss=:mse, nrounds=20, lr=1e-2) m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names) p = m(deval) @@ -39,4 +42,44 @@ end +end + +@testset "Embeddings - Classification" begin + + Random.seed!(123) + X, y = @load_crabs + df = DataFrame(X) + df[!, :class] = y + target_name = "class" + feature_names = setdiff(names(df), [target_name]) + + train_indices = randperm(nrow(df))[1:Int(0.8 * nrow(df))] + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + @testset "$et" for et in [:periodic, :linear, :piecewise] + + + embed_kw = Dict(:embedding_type => et, :d_embedding => 8) + if et == :piecewise + embed_kw[:n_bins] = 16 + end + if et == :periodic + embed_kw[:n_frequencies] = 8 + end + + arch = NeuroTabModels.TabMConfig(; k=2, n_blocks=2, d_block=16, dropout=0.0, scaling_init=:random_signs) + learner = NeuroTabClassifier(arch; + nrounds=500, batchsize=32, early_stopping_rounds=100, lr=1e-2, + embedding_config=embed_kw) + + m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names) + + ptrain = [argmax(x) for x in eachrow(m(dtrain))] + peval = [argmax(x) for x in eachrow(m(deval))] + @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 + @test mean(peval .== levelcode.(deval.class)) > 0.95 + + end + end \ No newline at end of file From 9c73c9b675d5015ad52e83b1fc28fea33740e23f Mon Sep 17 00:00:00 2001 From: Jeremie Desgagne-Bouchard Date: Tue, 14 Apr 2026 16:46:24 -0400 Subject: [PATCH 087/120] Embed (#34) * - Add EmbeddingConfig as architecture-agnostic embedding configuration - Move embedding construction from TabMConfig to fit.jl (_build_chain) - Embeddings now specified via embedding_config kwarg on learner constructors - Support Dict -> EmbeddingConfig conversion for ergonomic API - Add activation function parameter to LinearEmbeddings, PeriodicEmbeddings, PiecewiseLinearEmbeddings (replaces hardcoded relu/bool) - Fix PiecewiseLinearEncoding to use clamp(0,1) per paper specification - Unify _ensure_3d/_bcast into _reshape_3d in losses.jl - Rename _activation to _inverse_link in infer.jl - Remove LinearReLUEmbeddings (redundant with LinearEmbeddings(activation=relu)) - Add embedding regression + classification tests Signed-off-by: AdityaPandeyCN * tune test params Signed-off-by: AdityaPandeyCN * Simplify chain construction and restore scaling_init override Signed-off-by: AdityaPandeyCN * resolve scaling_init issue Signed-off-by: AdityaPandeyCN * up * up * k support for neurotrees * fix oblivious * api naming * config * fix neuro statck * typo * Embed scalers (#38) * up * MOETrees regression scalers * scalers * Sync (#37) * up * deps * up * up * up * revert relu fix MOE * fix Reactant version --------- Signed-off-by: AdityaPandeyCN Co-authored-by: AdityaPandeyCN Co-authored-by: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> --- Project.toml | 8 +- benchmarks/YEAR-regression.jl | 63 +++++---- benchmarks/benchmark_mse.jl | 1 - benchmarks/titanic-logloss.jl | 61 +++++---- experiments/attention.jl | 76 +++++++++++ experiments/tree-tensor.jl | 70 +++++++--- src/Fit/callback.jl | 15 ++- src/Fit/fit.jl | 55 +++++--- src/data.jl | 123 +++++++++++++---- src/infer.jl | 86 ++++++++---- src/models/MOETree/model.jl | 115 ++++++++++++++++ src/models/MOETree/moetrees.jl | 156 ++++++++++++++++++++++ src/models/NeuroTree/model.jl | 37 ++--- src/models/NeuroTree/neurotrees.jl | 28 ++-- src/models/embeddings/batchnorm.jl | 26 ++++ src/models/embeddings/compute_bins.jl | 26 ++-- src/models/embeddings/config.jl | 45 ++++--- src/models/embeddings/embeddings.jl | 1 + src/models/embeddings/linear.jl | 20 +-- src/models/embeddings/nlinear.jl | 2 +- src/models/embeddings/periodic.jl | 32 ++--- src/models/embeddings/piecewise_linear.jl | 28 ++-- src/models/models.jl | 5 +- 23 files changed, 827 insertions(+), 252 deletions(-) create mode 100644 experiments/attention.jl create mode 100644 src/models/MOETree/model.jl create mode 100644 src/models/MOETree/moetrees.jl create mode 100644 src/models/embeddings/batchnorm.jl diff --git a/Project.toml b/Project.toml index 092d19c..4b53e94 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NeuroTabModels" uuid = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" -authors = ["jeremie.db "] version = "0.4.0" +authors = ["jeremie.db "] [deps] CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" @@ -24,13 +24,17 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" [compat] CategoricalArrays = "1" DataFrames = "1.3" +Enzyme = "0.13" Lux = "1" +LuxCore = "1" +LuxLib = "1" MLJModelInterface = "1.2.1" MLUtils = "0.4" NNlib = "0.9" +OneHotArrays = "0.2" Optimisers = "0.4" Random = "1" -Reactant = "0.2.236" +Reactant = "=0.2.247" Statistics = "1" StatsBase = "0.34" Tables = "1.9" diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 200baac..094001c 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -2,6 +2,7 @@ using Random using CSV using DataFrames using Statistics: mean, std +using StatsBase: tiedrank using NeuroTabModels using AWS: AWSCredentials, AWSConfig, @service @service S3 @@ -22,18 +23,14 @@ path = "share/data/year/year-eval-idx.txt" raw = S3.get_object("jeremiedb", path, Dict("response-content-type" => "application/octet-stream"); aws_config) eval_idx = DataFrame(CSV.File(raw, header=false))[:, 1] .+ 1 -transform!(df_tot, "Column1" => identity => "y_raw") -select!(df_tot, Not("Column1")) -transform!(df_tot, "y_raw" => (x -> (x .- mean(x)) ./ std(x)) => "y_norm") -# transform!(df_tot, "y_raw" => (x -> (x .- minimum(x)) ./ std(x)) => "y_norm") -feature_names = setdiff(names(df_tot), ["y_raw", "y_norm", "w"]) +target_name = "y" +rename!(df_tot, "Column1" => target_name) +feature_names = setdiff(names(df_tot), ["y", "w"]) df_tot.w .= 1.0 -target_name = "y_norm" # function percent_rank(x::AbstractVector{T}) where {T} # return tiedrank(x) / (length(x) + 1) # end - # transform!(df_tot, feature_names .=> percent_rank .=> feature_names) dtrain = df_tot[train_idx, :]; @@ -42,27 +39,31 @@ dtest = df_tot[(end-51630+1):end, :]; arch = NeuroTabModels.NeuroTreeConfig(; tree_type=:binary, - proj_size=4, actA=:identity, - depth=4, + k=1, ntrees=32, + depth=4, stack_size=1, - hidden_size=1, - init_scale=0.0, + hidden_size=16, + init_scale=0.1, scaler=true, - MLE_tree_split=false ) + +# arch = NeuroTabModels.MOETreeConfig(; +# tree_type=:binary, +# depth=5, +# ntrees=8, +# stack_size=1, +# init_scale=0.1, +# ) + # arch = NeuroTabModels.TabMConfig(; # arch_type=:tabm, -# k=32, -# d_block=128, +# k=16, +# d_block=64, # n_blocks=3, # dropout=0.1, -# n_bins=16, -# use_embeddings=true, -# embedding_type=:linear, # periodic, piecewise, linear -# d_embedding=8, -# scaling_init=:normal, +# # scaling_init=:normal, # ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, @@ -78,33 +79,43 @@ arch = NeuroTabModels.NeuroTreeConfig(; # ) device = :gpu -loss = :gaussian_mle # :mse :gaussian_mle :tweedie +loss = :mse # :mse :gaussian_mle :tweedie + +# embedding_config = Dict( +# :embedding_type => :piecewise, +# :d_embedding => 8, +# :activation => nothing, +# :bins => 16, +# :frequencies => 16, +# ) +embedding_config = Dict(:embedding_type => :batchnorm) learner = NeuroTabRegressor( arch; + embedding_config, loss, nrounds=200, early_stopping_rounds=2, - lr=3e-4, - batchsize=1024, + lr=1e-3, + batchsize=512, device ) -m = NeuroTabModels.fit( +@time m = NeuroTabModels.fit( learner, dtrain; deval, target_name, feature_names, print_every_n=5, -) +); p_eval = m(deval; device=:cpu); p_eval = p_eval[:, 1] -mse_eval = mean((p_eval .- deval.y_norm) .^ 2) +mse_eval = mean((p_eval .- deval.y) .^ 2) @info "MSE - deval" mse_eval p_test = m(dtest; device=:cpu); p_test = p_test[:, 1] -mse_test = mean((p_test .- dtest.y_norm) .^ 2) * std(df_tot.y_raw)^2 +mse_test = mean((p_test .- dtest.y) .^ 2) @info "MSE - dtest" mse_test diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index f5a6b14..28b3772 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -18,7 +18,6 @@ target_name = "y" # arch = NeuroTabModels.NeuroTreeConfig(; # tree_type=:binary, -# proj_size=1, # actA=:identity, # init_scale=1.0, # depth=4, diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 4137b7a..69366cc 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -32,37 +32,53 @@ deval = df[setdiff(1:nrow(df), train_indices), :] target_name = "Survived" feature_names = setdiff(names(df), ["Survived"]) -# arch = NeuroTabModels.NeuroTreeConfig(; +arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + k=8, + depth=4, + ntrees=16, + stack_size=2, + hidden_size=8, + actA=:identity, + init_scale=1.0, + scaler=true, +) + +# arch = NeuroTabModels.MOETreeConfig(; # tree_type=:binary, -# proj_size=1, -# init_scale=1.0, -# depth=4, -# ntrees=16, +# k=1, +# depth=3, +# ntrees=8, # stack_size=1, -# hidden_size=1, -# actA=:identity, -# MLE_tree_split=false, +# hidden_size=8, +# init_scale=1.0, # ) -arch = NeuroTabModels.TabMConfig(; - arch_type=:tabm, - k=32, - d_block=64, - n_blocks=3, - dropout=0.1, - bins=nothing, - use_embeddings=false, - embedding_type=:periodic, - d_embedding=16, - scaling_init=:random_signs, -) +# arch = NeuroTabModels.TabMConfig(; +# arch_type=:tabm, +# k=8, +# d_block=32, +# n_blocks=3, +# dropout=0.1, +# scaling_init=:normal, +# ) + +# embedding_config = Dict( +# :embedding_type => :piecewise, +# :d_embedding => 8, +# :activation => nothing, +# :bins => 16, +# :frequencies => 16, +# ) +embedding_config = Dict(:embedding_type => :batchnorm) learner = NeuroTabRegressor( arch; + embedding_config, loss=:logloss, - nrounds=100, + nrounds=200, early_stopping_rounds=2, lr=1e-2, - device=:cpu + device=:gpu ) @time m = NeuroTabModels.fit( @@ -74,7 +90,6 @@ learner = NeuroTabRegressor( print_every_n=10, ); -# commented out - inference not yet adapted p_train = m(dtrain) p_eval = m(deval) @info mean((p_train .> 0.5) .== (dtrain[!, target_name] .> 0.5)) diff --git a/experiments/attention.jl b/experiments/attention.jl new file mode 100644 index 0000000..d57ac79 --- /dev/null +++ b/experiments/attention.jl @@ -0,0 +1,76 @@ +using NeuroTabModels +using DataFrames +using BenchmarkTools +using Dates +using Random: seed! +import Base: length, getindex +import MLUtils: DataLoader +using Lux +using Statistics + +Threads.nthreads() + +seed!(123) +nobs = Int(1e4) +num_feat = Int(10) +X = rand(Float32, nobs, num_feat) +Y = randn(Float32, size(X, 1)) +dtrain = DataFrame(X, :auto) +feature_names = names(dtrain) +dtrain.y = Y +target_name = "y" +dtrain.w = rand(nrow(dtrain)) +weight_name = "w" +dtrain.date = rand(Date("2026-01-01"):Date("2026-01-05"), nrow(dtrain)) + +arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + actA=:identity, + init_scale=1.0, + depth=4, + ntrees=16, + stack_size=1, + hidden_size=1, + scaler=true, +) + +learner = NeuroTabRegressor( + arch; + loss=:mse, + nrounds=10, + lr=1e-2, + batchsize=2048, + device=:gpu +) + +sort!(dtrain, :date) +m = NeuroTabModels.fit( + learner, + dtrain; + deval=dtrain, + target_name, + feature_names, + weight_name, + print_every_n=2, +); +p_train_1 = m(dtrain; device=:gpu); +p_train_2 = m(dtrain; device=:gpu); + +m = NeuroTabModels.fit( + learner, + dtrain; + deval=dtrain, + target_name, + feature_names, + weight_name, + group_key=:date, + print_every_n=2, +); +p_train_grp = m(dtrain; device=:gpu); + +cor(p_train_1, p_train_2) +cor(p_train_1, p_train_grp) + +dfg = groupby(dtrain, :date) +data = NeuroTabModels.Data.get_df_loader_train(dfg; feature_names, target_name, weight_name) +length(data) diff --git a/experiments/tree-tensor.jl b/experiments/tree-tensor.jl index 9cb7960..477a216 100644 --- a/experiments/tree-tensor.jl +++ b/experiments/tree-tensor.jl @@ -1,10 +1,9 @@ - using BenchmarkTools using Random using NeuroTabModels using NeuroTabModels.Models.NeuroTrees -using CUDA -# using CairoMakie +# using CUDA +using CairoMakie # density(m.chain.layers[2].trees[1].b) # density(m.chain.layers[2].trees[1].s) @@ -15,16 +14,16 @@ using CUDA # mean(abs.(vec(m.chain.layers[2].trees[1].w)) .< 1e-1) ############################### -# original mask +# softplus mask ############################### -mask = NeuroTrees.get_mask(Val(:binary), 4) -mask = NeuroTrees.get_mask(Val(:oblivious), 4) +mask = NeuroTrees.get_softplus_mask(Val(:binary), 4) +mask = NeuroTrees.get_softplus_mask(Val(:oblivious), 4) fig = Figure(; size=(450, 450)) ax = Axis(fig[1, 1]; - title="original mask", - xlabel="Nodes", - ylabel="Leaves", + title="softplus mask", + xlabel="Leaves", + ylabel="Nodes", aspect=DataAspect(), xticks=collect(1:size(mask, 1)), yticks=collect(1:size(mask, 2)), @@ -35,13 +34,8 @@ hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, translate!(hm, 0, 0, -100) fig -############################### -# softplus mask -############################### +fig = Figure(; size=(900, 450)) mask = NeuroTrees.get_softplus_mask(Val(:binary), 4) -mask = NeuroTrees.get_softplus_mask(Val(:oblivious), 4) - -fig = Figure(; size=(450, 450)) ax = Axis(fig[1, 1]; title="softplus mask", xlabel="Leaves", @@ -54,6 +48,20 @@ ax = Axis(fig[1, 1]; ) hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) translate!(hm, 0, 0, -100) + +mask = NeuroTrees.get_softplus_mask(Val(:oblivious), 4) +ax = Axis(fig[1, 2]; + title="softplus mask", + xlabel="Leaves", + ylabel="Nodes", + aspect=DataAspect(), + xticks=collect(1:size(mask, 1)), + yticks=collect(1:size(mask, 2)), + xgridcolor=:lightgrey, + ygridcolor=:lightgrey +) +hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) +translate!(hm, 0, 0, -100) fig ############################### @@ -77,6 +85,38 @@ hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, translate!(hm, 0, 0, -100) fig + +fig = Figure(; size=(900, 450)) +mask = NeuroTrees.get_logits_mask(Val(:binary), 4) +ax = Axis(fig[1, 1]; + title="logits mask", + xlabel="Leaves", + ylabel="Nodes", + aspect=DataAspect(), + xticks=collect(1:size(mask, 1)), + yticks=collect(1:size(mask, 2)), + xgridcolor=:lightgrey, + ygridcolor=:lightgrey +) +hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) +translate!(hm, 0, 0, -100) + +mask = NeuroTrees.get_logits_mask(Val(:oblivious), 4) +ax = Axis(fig[1, 2]; + title="logits mask", + xlabel="Leaves", + ylabel="Nodes", + aspect=DataAspect(), + xticks=collect(1:size(mask, 1)), + yticks=collect(1:size(mask, 2)), + xgridcolor=:lightgrey, + ygridcolor=:lightgrey +) +hm = heatmap!(ax, 1:size(mask, 1)+1, 1:size(mask, 2)+1, mask, colormap=[:white, "#1f4e79"]) +translate!(hm, 0, 0, -100) +fig + + function leaf_weights(logits, lmask, smask) logits * lmask .+ softplus.(logits) * smask end diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 11b1124..e16de91 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -9,7 +9,7 @@ using ..Data: get_df_loader_train using ..Metrics using Lux: Training, reactant_device, testmode -using Reactant: @compile +using Reactant: @compile, set_default_backend export CallBack, init_logger, update_logger!, agg_logger @@ -26,17 +26,22 @@ end function CallBack( config::LearnerTypes, - deval::AbstractDataFrame, - ts::Training.TrainState; + df::AbstractDataFrame, + cache; feature_names, target_name, weight_name=nothing, - offset_name=nothing + offset_name=nothing, + group_key=nothing ) dev = reactant_device() + ts = cache[:train_state] + scalers = cache[:scalers] batchsize = config.batchsize feval = metric_dict[config.metric] - deval = get_df_loader_train(deval; feature_names, target_name, weight_name, offset_name, batchsize, shuffle=false) |> dev + + dfg = isnothing(group_key) ? df : groupby(df, group_key; sort=true) + deval = get_df_loader_train(dfg; feature_names, target_name, weight_name, offset_name, scalers, batchsize, shuffle=false) |> dev ps, st = ts.parameters, testmode(ts.states) d0 = first(deval) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index b86c54d..63cb023 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -10,9 +10,9 @@ using ..Metrics using ..Infer: reduce_pred import Random: Xoshiro +import Statistics: mean, std import MLJModelInterface: fit import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam - using Lux using Reactant using Lux: cpu_device, reactant_device @@ -35,17 +35,25 @@ function init( feature_names, target_name, weight_name=nothing, - offset_name=nothing + offset_name=nothing, + group_key=nothing ) + + feature_names, target_name = Symbol.(feature_names), Symbol(target_name) + weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) + offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) + group_key = isnothing(group_key) ? nothing : Symbol(group_key) + dev = _get_device(config) batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) lux_loss = get_loss_fn(L) + outsize = 1 + scalers = nothing target_levels = nothing target_isordered = false - outsize = 1 if L <: MLogLoss eltype(df[!, target_name]) <: CategoricalValue || error("Target `$target_name` must be `<: CategoricalValue`") @@ -54,17 +62,13 @@ function init( outsize = length(target_levels) elseif L <: GaussianMLE outsize = 2 + scalers = (mu=mean(df[!, target_name]), sigma=std(df[!, target_name])) + elseif L <: Union{MSE,MAE} + scalers = (mu=mean(df[!, target_name]), sigma=std(df[!, target_name])) end - data = get_df_loader_train(df; feature_names, target_name, weight_name, offset_name, batchsize) |> dev - - info = Dict( - :nrounds => 0, - :feature_names => feature_names, - :target_levels => target_levels, - :target_isordered => target_isordered, - :device => config.device - ) + dfg = isnothing(group_key) ? df : groupby(df, group_key; sort=true) + data = get_df_loader_train(dfg; feature_names, target_name, weight_name, offset_name, scalers, batchsize) |> dev # Build chain: optional embeddings + architecture backbone embed_config = config.embedding_config @@ -72,16 +76,27 @@ function init( chain = config.arch(; nfeats, outsize) else if embed_config.embedding_type == :piecewise - X_train = Matrix{Float32}(df[:, feature_names]) + x_train = Matrix{Float32}(df[:, feature_names]) else - X_train = nothing + x_train = nothing end - embed_chain = embed_config(; nfeats, X_train) + embed_chain = embed_config(; nfeats, x_train) d_in = nfeats * embed_config.d_embedding d_features = fill(embed_config.d_embedding, nfeats) chain = Chain(embed_chain, config.arch(; nfeats=d_in, outsize, d_features, scaling_init_override=:normal)) end + info = Dict( + :nrounds => 0, + :feature_names => feature_names, + :target_name => target_name, + :weight_name => weight_name, + :offset_name => offset_name, + :group_key => group_key, + :target_levels => target_levels, + :target_isordered => target_isordered, + :scalers => scalers + ) m = NeuroTabModel(L, chain, info) rng = Xoshiro(config.seed) @@ -89,7 +104,7 @@ function init( opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) - return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts) + return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts, :scalers => scalers) end """ @@ -135,19 +150,17 @@ function fit( target_name, weight_name=nothing, offset_name=nothing, + group_key=nothing, deval=nothing, print_every_n=9999, verbosity=1 ) - feature_names, target_name = Symbol.(feature_names), Symbol(target_name) - weight_name = isnothing(weight_name) ? nothing : Symbol(weight_name) - offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) - m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name) + m, cache = init(config, dtrain; feature_names, target_name, weight_name, offset_name, group_key) logger = nothing if !isnothing(deval) - cb = CallBack(config, deval, cache[:train_state]; feature_names, target_name, weight_name, offset_name) + cb = CallBack(config, deval, cache; feature_names, target_name, weight_name, offset_name) logger = init_logger(config) cb(logger, 0, cache[:train_state]) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] diff --git a/src/data.jl b/src/data.jl index a6ee8df..cb6dba2 100644 --- a/src/data.jl +++ b/src/data.jl @@ -11,7 +11,7 @@ using CategoricalArrays """ ContainerTrain """ -struct ContainerTrain{A<:AbstractMatrix,B,C,D} +struct ContainerTrain{A,B,C,D} x::A y::B w::C @@ -19,6 +19,7 @@ struct ContainerTrain{A<:AbstractMatrix,B,C,D} end length(data::ContainerTrain) = size(data.x, 2) +length(data::ContainerTrain{<:Vector}) = length(data.x) function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B,C<:Nothing,D<:Nothing} x = data.x[:, idx] @@ -46,6 +47,14 @@ function getindex(data::ContainerTrain{A,B,C,D}, idx::AbstractVector) where {A,B return (x, y, w, offset) end +# for GroupedDataFrame +function getindex(data::ContainerTrain{A,B,C,D}, idx::Int) where {A<:Vector,B,C<:Vector,D<:Nothing} + x = data.x[idx] + y = data.y[idx] + w = data.w[idx] + return (x, y, w) +end + function get_df_loader_train( df::AbstractDataFrame; feature_names, @@ -53,6 +62,7 @@ function get_df_loader_train( weight_name=nothing, offset_name=nothing, batchsize, + scalers=nothing, shuffle=true) feature_names = Symbol.(feature_names) @@ -63,8 +73,11 @@ function get_df_loader_train( else y = Float32.(df[!, target_name]) end - y = reshape(y, 1, :) + if !isnothing(scalers) + y .= (y .- scalers[:mu]) ./ scalers[:sigma] + end + y = reshape(y, 1, :) w = isnothing(weight_name) ? nothing : Float32.(df[!, weight_name]) offset = if isnothing(offset_name) @@ -79,50 +92,108 @@ function get_df_loader_train( return dtrain end +function get_df_loader_train( + dfg::GroupedDataFrame; + feature_names, + target_name, + weight_name=nothing, + offset_name=nothing, + batchsize=0, + scalers=nothing, + shuffle=true) + + n = length(dfg) + nfeats = length(feature_names) + bs = maximum(dfg.ends .- dfg.starts) + 1 + @info "groupedDF loader" bs + + x = [zeros(Float32, nfeats, bs) for _ in 1:n] + y = [zeros(Float32, bs) for _ in 1:n] + w = [zeros(Float32, bs) for _ in 1:n] + + for i in 1:n + df = dfg[i] + x[i][:, 1:nrow(df)] .= Matrix(df[!, feature_names])' + if isnothing(scalers) + y[i][1:nrow(df)] .= df[!, target_name] + else + y[i][1:nrow(df)] .= (df[!, target_name] .- scalers[:mu]) ./ scalers[:sigma] + end + w[i][1:nrow(df)] .= 1.0 + end + offset = nothing + + container = ContainerTrain(x, y, w, offset) + dtrain = DataLoader(container; shuffle, batchsize=0, partial=false, parallel=false) + return dtrain +end + + """ ContainerInfer """ -struct ContainerInfer{A<:AbstractMatrix,D} +struct ContainerInfer{A<:AbstractMatrix} x::A - offset::D end - length(data::ContainerInfer) = size(data.x, 2) -function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:Nothing} +function getindex(data::ContainerInfer, idx::AbstractVector) x = data.x[:, idx] return x end -function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:AbstractVector} - x = data.x[:, idx] - offset = data.offset[idx] - return (x, offset) -end -function getindex(data::ContainerInfer{A,D}, idx::AbstractVector) where {A,D<:AbstractMatrix} - x = data.x[:, idx] - offset = data.offset[:, idx] - return (x, offset) -end function get_df_loader_infer( df::AbstractDataFrame; feature_names, - offset_name=nothing, - batchsize) - + batchsize +) feature_names = Symbol.(feature_names) x = Matrix{Float32}(Matrix{Float32}(select(df, feature_names))') - offset = if isnothing(offset_name) - nothing - else - isa(offset_name, String) ? Float32.(df[!, offset_name]) : Matrix{Float32}(Matrix{Float32}(df[!, offset_name])') - end - - container = ContainerInfer(x, offset) + container = ContainerInfer(x) batchsize = min(batchsize, length(container)) dinfer = DataLoader(container; shuffle=false, batchsize, partial=true, parallel=false) return dinfer end +""" + ContainerInferGrp +""" +struct ContainerInferGrp{A<:AbstractVector,B<:AbstractVector} + x::A + mask::B +end +length(data::ContainerInferGrp) = length(data.x) + +# for GroupedDataFrame +function getindex(data::ContainerInferGrp, idx::Int) + x = data.x[idx] + mask = data.mask[idx] + return (x, mask) +end + +function get_df_loader_infer( + dfg::GroupedDataFrame; + feature_names, + batchsize=0 +) + n = length(dfg) + nfeats = length(feature_names) + bs = maximum(dfg.ends .- dfg.starts) + 1 + @info "groupedDF loader" bs + + x = [zeros(Float32, nfeats, bs) for _ in 1:n] + mask = [zeros(Bool, bs) for _ in 1:n] + + for i in 1:n + df = dfg[i] + x[i][:, 1:nrow(df)] .= Matrix(df[!, feature_names])' + mask[i][1:nrow(df)] .= true + end + + container = ContainerInferGrp(x, mask) + dinfer = DataLoader(container; shuffle=false, batchsize=0, partial=false, parallel=false) + return dinfer +end + end #module \ No newline at end of file diff --git a/src/infer.jl b/src/infer.jl index f1b9745..9b39654 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -10,7 +10,7 @@ using Reactant using Reactant: @compile using NNlib: sigmoid, softmax using Statistics: mean -using DataFrames: AbstractDataFrame +using DataFrames: AbstractDataFrame, GroupedDataFrame, groupby import MLUtils: DataLoader export infer, reduce_pred @@ -35,62 +35,90 @@ _assemble(::Type{<:GaussianMLE}, raw_preds) = reduce(hcat, raw_preds) _assemble(::Type, raw_preds) = vcat([vec(p) for p in raw_preds]...) # Apply inverse link to convert from model scale to natural scale -_inverse_link(::Type{<:LogLoss}, pred) = sigmoid.(pred) -_inverse_link(::Type{<:Tweedie}, pred) = exp.(pred) -_inverse_link(::Type{<:Union{MSE,MAE}}, pred) = pred +_inverse_link(::Type{<:LogLoss}, pred, scalers) = sigmoid.(pred) +_inverse_link(::Type{<:Tweedie}, pred, scalers) = exp.(pred) +_inverse_link(::Type{<:Union{MSE,MAE}}, pred, scalers) = pred .* scalers[:sigma] .+ scalers[:mu] -function _inverse_link(::Type{<:MLogLoss}, pred) +function _inverse_link(::Type{<:MLogLoss}, pred, scalers) return Matrix(softmax(pred; dims=1)') end -function _inverse_link(::Type{<:GaussianMLE}, pred) +function _inverse_link(::Type{<:GaussianMLE}, pred, scalers) p = Matrix(pred') - p[:, 2] .= exp.(p[:, 2]) + @views p[:, 1] .= p[:, 1] .* scalers[:sigma] .+ scalers[:mu] + @views p[:, 2] .= exp.(p[:, 2]) .* scalers[:sigma] return p end -function _postprocess(::Type{L}, raw_preds; raw::Bool=false) where {L} - assembled = _assemble(L, raw_preds) - raw && return assembled - return _inverse_link(L, assembled) -end - -function infer(m::NeuroTabModel{L}, data; device=:cpu, raw::Bool=false) where {L} +function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} dev = _get_device(device) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) + scalers = m.info[:scalers] - raw_preds = Vector{AbstractArray}() - - b_first = first(data) - x0 = b_first isa Tuple ? b_first[1] : b_first + x0 = first(data) compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) - for b in data - x = b isa Tuple ? b[1] : b + preds = Vector{AbstractArray}() + for x in data if size(x) == size(x0) pred = compiled(m.chain, dev(x), ps, st) else pred = Reactant.@jit _forward_reduce(m.chain, dev(x), ps, st) end - push!(raw_preds, cdev(pred)) + push!(preds, cdev(pred)) end - return _postprocess(L, raw_preds; raw) + p_raw = _assemble(L, preds) + proj || return p_raw + return _inverse_link(L, p_raw, scalers) end -function infer(m::NeuroTabModel, data::AbstractDataFrame; device=:cpu, raw::Bool=false) - dinfer = get_df_loader_infer(data; feature_names=m.info[:feature_names], batchsize=2048) - return infer(m, dinfer; device, raw) +function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} + dev = _get_device(device) + cdev = cpu_device() + ps = dev(m.info[:ps]) + st = dev(m.info[:st]) + scalers = m.info[:scalers] + + (x0, mask0) = first(data) + # @info typeof("mask0") mask0 + # data = data |> dev + # (x0, mask0) = first(data) + # @info typeof("mask0-dev") mask0 + compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) + + preds = Vector{AbstractArray}() + for (x, mask) in data + pred = compiled(m.chain, dev(x), ps, st) + push!(preds, cdev(pred)[mask]) + end + + p_raw = _assemble(L, preds) + proj || return p_raw + return _inverse_link(L, p_raw, scalers) end -function (m::NeuroTabModel)(data::AbstractDataFrame; device=:cpu) - return infer(m, data; device) +function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, proj::Bool=true) + group_key = m.info[:group_key] + if isnothing(group_key) + dinfer = get_df_loader_infer(df; feature_names=m.info[:feature_names], batchsize=2048) + p = infer(m, dinfer; device, proj) + else + dfg = groupby(df, group_key; sort=true) + dinfer = get_df_loader_infer(dfg; feature_names=m.info[:feature_names], batchsize=2048) + p = infer_grp(m, dinfer; device, proj) + end + return p end -function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) - return infer(m, [(x,)]; device) +function (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, proj::Bool=true) + return infer(m, df; device, proj) end +# function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) +# return infer(m, [(x,)]; device) +# end + end \ No newline at end of file diff --git a/src/models/MOETree/model.jl b/src/models/MOETree/model.jl new file mode 100644 index 0000000..36c0178 --- /dev/null +++ b/src/models/MOETree/model.jl @@ -0,0 +1,115 @@ +struct MOETree{F} <: AbstractLuxLayer + tree_type::Symbol + actA::F + scaler::Bool + feats::Int + outs::Int + depth::Int + trees::Int + nodes::Int + leaves::Int + k::Int + init_scale::Float32 +end + +function MOETree((feats, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, actA=identity, scaler=true, depth, trees, k, init_scale=0.1) + @assert tree_type ∈ [:binary, :oblivious] + nodes = tree_type == :binary ? 2^depth - 1 : depth + leaves = 2^depth + return MOETree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, k, Float32(init_scale)) +end + +# Define the Lux interface +function LuxCore.initialparameters(rng::AbstractRNG, l::MOETree) + return ( + # router + lw=Float32.((rand(rng, l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # [NTK,F] + lb=zeros(Float32, l.nodes * l.trees), # [NTK] + ls=Float32.(fill(log(expm1(1)), l.nodes * l.trees)), # [NTK] + # expert-model + w=Float32.((rand(rng, l.nodes * l.trees * l.leaves, l.feats) .- 0.5) ./ 4), # [NTK,F] + b=zeros(Float32, l.nodes * l.trees * l.leaves), # [NTK] + s=Float32.(fill(log(expm1(1)), l.nodes * l.trees * l.leaves)), # [NTK] + p=Float32.(randn(rng, l.outs, l.leaves, l.trees, l.leaves) .* l.init_scale), # [P,L,T,K] + ) +end +function LuxCore.initialstates(rng::AbstractRNG, l::MOETree) + return ( + ml=Float32.(get_logits_mask(Val(l.tree_type), l.depth)), + ms=Float32.(get_softplus_mask(Val(l.tree_type), l.depth)), + ) +end + +function (l::MOETree)(x, ps, st) + + enw1 = softplus(ps.ls) .* (ps.lw * x .+ ps.lb) # [F,B] => [NT,B] + enw = reshape(enw1, size(st.ml, 2), :) # [NTK,B] => [N,TB] + lwe1 = exp.(st.ml * enw .- st.ms * softplus.(enw)) # [N,TB] => [L,TB] + lwe2 = reshape(lwe1, 1, l.leaves, l.trees, size(x, 2)) # [L,TB] => [1,L,T,B] + lwe = dropdims(mean(lwe2; dims=3); dims=3) # [1,L,T,B] => [1,L,B] + + nw1 = softplus(ps.s) .* (ps.w * x .+ ps.b) # [F,B] => [NTL,B] + nw = reshape(nw1, size(st.ml, 2), :) # [NTK,B] => [N,TLB] + lw1 = exp.(st.ml * nw .- st.ms * softplus.(nw)) # [N,TLB] => [L,TLB] + lw = reshape(lw1, 1, l.leaves, l.trees, l.leaves, size(x, 2)) # [L,TLB] => [1,L,T,L,B] + y1 = dropdims(sum(ps.p .* lw; dims=2); dims=2) # [P,L,T,L] .* [1,L,T,L,B] => [P,T,L,B] + y2 = dropdims(mean(y1; dims=2); dims=2) # [P,T,L,B] => [P,L,B] + + y = sum(lwe .* y2; dims=2) # [1,L,B] .* [P,L,B] => [P,1,B] + + return y, st +end + +""" + get_logits_mask(::Val{:binary}, depth::Integer) +""" +function get_logits_mask(::Val{:binary}, depth::Integer) + nodes = 2^depth - 1 + leaves = 2^depth + mask = zeros(Bool, leaves, nodes) + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d) + stride = 2 * k + for b in 1:blocks + view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= 1 + end + end + return mask +end +function get_logits_mask(::Val{:oblivious}, depth::Integer) + leaves = 2^depth + mask = zeros(Bool, leaves, depth) + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d) + stride = 2 * k + for b in 1:blocks + view(mask, (b-1)*stride+1:(b-1)*stride+k, d) .= 1 + end + end + return mask +end + +""" + get_softplus_mask(::Val{:binary}, depth::Integer) +""" +function get_softplus_mask(::Val{:binary}, depth::Integer) + nodes = 2^depth - 1 + leaves = 2^depth + mask = zeros(Bool, leaves, nodes) + for d in 1:depth + blocks = 2^(d - 1) + k = 2^(depth - d + 1) + stride = k + for b in 1:blocks + view(mask, (b-1)*stride+1:(b-1)*stride+k, 2^(d - 1) + b - 1) .= 1 + end + end + return mask +end +function get_softplus_mask(::Val{:oblivious}, depth::Integer) + leaves = 2^depth + mask = ones(Bool, leaves, depth) + return mask +end diff --git a/src/models/MOETree/moetrees.jl b/src/models/MOETree/moetrees.jl new file mode 100644 index 0000000..b1d243d --- /dev/null +++ b/src/models/MOETree/moetrees.jl @@ -0,0 +1,156 @@ +module MOETrees + +export MOETreeConfig + +using Random +using Lux +using LuxCore +using Statistics: mean +using NNlib: softplus, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh, tanhshrink + +import ..Losses: get_loss_type, GaussianMLE +import ..Models: Architecture + +include("model.jl") + +struct StackedMOETree{L} <: LuxCore.AbstractLuxWrapperLayer{:chain} + chain::L +end + +function StackedMOETree((ins, outs)::Pair{<:Integer,<:Integer}; hidden_size::Int, stack_size::Int, k::Int=1, tree_kwargs...) + if stack_size == 1 + return StackedMOETree(MOETree(ins => outs; k, tree_kwargs...)) + end + + layers = Any[MOETree(ins => 1; k=hidden_size, tree_kwargs...), FlattenLayer()] + for _ in 2:(stack_size-1) + push!(layers, SkipConnection( + Chain(MOETree(hidden_size => 1; k=hidden_size, tree_kwargs...), FlattenLayer()), + + )) + end + push!(layers, MOETree(hidden_size => outs; k, tree_kwargs...)) + + return StackedMOETree(Chain(layers...)) +end + +struct MOETreeConfig <: Architecture + tree_type::Symbol + actA::Symbol + depth::Int + ntrees::Int + k::Int + hidden_size::Int + stack_size::Int + scaler::Bool + init_scale::Float32 + MLE_tree_split::Bool +end + +function MOETreeConfig(; kwargs...) + args = Dict{Symbol,Any}( + :tree_type => :binary, + :actA => :identity, + :depth => 4, + :ntrees => 32, + :k => 1, + :hidden_size => 1, + :stack_size => 1, + :scaler => true, + :init_scale => 0.1, + :MLE_tree_split => false, + ) + + args_ignored = setdiff(keys(kwargs), keys(args)) + length(args_ignored) > 0 && + @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(join(args_ignored, ", "))." + + args_default = setdiff(keys(args), keys(kwargs)) + length(args_default) > 0 && + @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." + + args_override = intersect(keys(args), keys(kwargs)) + for arg in args_override + args[arg] = kwargs[arg] + end + + return MOETreeConfig( + Symbol(args[:tree_type]), + Symbol(args[:actA]), + args[:depth], + args[:ntrees], + args[:k], + args[:hidden_size], + args[:stack_size], + args[:scaler], + args[:init_scale], + args[:MLE_tree_split], + ) +end + +function _tree_kwargs(config::MOETreeConfig) + return (; + config.tree_type, + config.depth, + trees=config.ntrees, + # k=config.k, + actA=act_dict[config.actA], + config.scaler, + config.init_scale, + ) +end + +function (config::MOETreeConfig)(; nfeats, outsize, kwargs...) + kwargs = _tree_kwargs(config) + + if config.MLE_tree_split + iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") + head_outsize = outsize ÷ 2 + chain = Chain( + Parallel( + vcat, + StackedMOETree(nfeats => head_outsize; config.hidden_size, config.stack_size, config.k, kwargs...), + StackedMOETree(nfeats => head_outsize; config.hidden_size, config.stack_size, config.k, kwargs...), + ), + ) + else + chain = Chain( + StackedMOETree(nfeats => outsize; config.hidden_size, config.stack_size, config.k, kwargs...), + ) + end + + return chain +end + +function _identity_act(x) + return x ./ sum(abs.(x), dims=2) +end +function _tanh_act(x) + x = tanh_fast.(x) + return x ./ sum(abs.(x), dims=2) +end +function _hardtanh_act(x) + x = hardtanh.(x) + return x ./ sum(abs.(x), dims=2) +end +function _tanhshrink_act(x) + x = tanhshrink.(x) + return x ./ sum(abs.(x), dims=2) +end + +""" + act_dict = Dict( + :identity => _identity_act, + :tanh => _tanh_act, + :hardtanh => _hardtanh_act, + :tanhshrink => _tanhshrink_act, + ) +Dictionary mapping features activation name to their function. +""" +const act_dict = Dict( + :identity => _identity_act, + :tanh => _tanh_act, + :hardtanh => _hardtanh_act, + :tanhshrink => _tanhshrink_act, +) + +end \ No newline at end of file diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 768748a..8b476b2 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -8,27 +8,30 @@ struct NeuroTree{F} <: AbstractLuxLayer trees::Int nodes::Int leaves::Int + k::Int init_scale::Float32 end -function NeuroTree(; feats, outs, tree_type=:binary, actA=identity, scaler=true, depth, trees, init_scale=0.1) - nodes = 2^depth - 1 +function NeuroTree(; feats, outs, tree_type=:binary, actA=identity, scaler=true, depth, trees, k, init_scale=0.1) + @assert tree_type ∈ [:binary, :oblivious] + nodes = tree_type == :binary ? 2^depth - 1 : depth leaves = 2^depth - return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, Float32(init_scale)) + return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, k, Float32(init_scale)) end -function NeuroTree((feats, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, actA=identity, scaler=true, depth, trees, init_scale=0.1) - nodes = 2^depth - 1 +function NeuroTree((feats, outs)::Pair{<:Integer,<:Integer}; tree_type=:binary, actA=identity, scaler=true, depth, trees, k, init_scale=0.1) + @assert tree_type ∈ [:binary, :oblivious] + nodes = tree_type == :binary ? 2^depth - 1 : depth leaves = 2^depth - return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, Float32(init_scale)) + return NeuroTree(tree_type, actA, scaler, feats, outs, depth, trees, nodes, leaves, k, Float32(init_scale)) end # Define the Lux interface function LuxCore.initialparameters(rng::AbstractRNG, l::NeuroTree) return ( - w=Float32.((rand(rng, l.nodes * l.trees, l.feats) .- 0.5) ./ 4), # w - b=zeros(Float32, l.nodes * l.trees), # b - s=Float32.(fill(log(exp(1) - 1), l.nodes * l.trees)), # s - p=Float32.(randn(rng, l.outs, l.leaves * l.trees) .* l.init_scale), # p + w=Float32.((rand(rng, l.nodes * l.trees * l.k, l.feats) .- 0.5) ./ 4), # [NTK,F] + b=zeros(Float32, l.nodes * l.trees * l.k), # [NTK] + s=Float32.(fill(log(expm1(1)), l.nodes * l.trees * l.k)), # [NTK] + p=Float32.(randn(rng, l.outs, l.leaves, l.trees) .* l.init_scale), # [P,L,T,K] ) end @@ -41,18 +44,18 @@ end function (l::NeuroTree)(x, ps, st) if l.scaler - nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] + nw = softplus(ps.s) .* (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NTK,B] else - nw = (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NT,B] + nw = (l.actA(ps.w) * x .+ ps.b) # [F,B] => [NTK,B] end - nw = reshape(nw, size(st.ml, 2), :) # [NT,B] => [N,TB] - lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) # [N,TB] => [L,TB] - lw = reshape(lw, :, size(x, 2)) # [L,TB] => [LT,B] - y = ps.p * lw ./ l.trees # [P,LT] * [LT,B] => [P,B] + nw = reshape(nw, size(st.ml, 2), :) # [NTK,B] => [N,TKB] + lw = exp.(st.ml * nw .- st.ms * softplus.(nw)) # [N,TKB] => [L,TKB] + lw = reshape(lw, 1, l.leaves, l.trees, l.k, size(x, 2)) # [L,TKB] => [1,L,T,K,B] + y1 = dropdims(sum(ps.p .* lw; dims=2); dims=2) # [P,L,T,K] * [1,L,T,K,B] => [P,T,K,B] + y = dropdims(mean(y1; dims=2); dims=2) # [P,T,K,B] => [P,K,B] return y, st end - """ get_logits_mask(::Val{:binary}, depth::Integer) """ diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index d4b9a35..4c262ca 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -5,6 +5,7 @@ export NeuroTreeConfig using Random using Lux using LuxCore +using Statistics: mean using NNlib: softplus, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh, tanhshrink import ..Losses: get_loss_type, GaussianMLE @@ -16,18 +17,18 @@ struct StackedNeuroTree{L} <: LuxCore.AbstractLuxWrapperLayer{:chain} chain::L end -function StackedNeuroTree(; feats::Int, outs::Int, hidden_size::Int, stack_size::Int, tree_kwargs...) +function StackedNeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; hidden_size::Int, stack_size::Int, k::Int=1, tree_kwargs...) if stack_size == 1 - return StackedNeuroTree(NeuroTree(; feats, outs, tree_kwargs...)) + return StackedNeuroTree(NeuroTree(ins => outs; k, tree_kwargs...)) end - layers = Any[NeuroTree(; feats, outs=hidden_size, tree_kwargs...)] - for _ in 1:(stack_size-2) + layers = Any[NeuroTree(ins => 1; k=hidden_size, tree_kwargs...), FlattenLayer()] + for _ in 2:(stack_size-1) push!(layers, SkipConnection( - NeuroTree(; feats=hidden_size, outs=hidden_size, tree_kwargs...), + + Chain(NeuroTree(hidden_size => 1; k=hidden_size, tree_kwargs...), FlattenLayer()), + )) end - push!(layers, NeuroTree(; feats=hidden_size, outs, tree_kwargs...)) + push!(layers, NeuroTree(hidden_size => outs; k, tree_kwargs...)) return StackedNeuroTree(Chain(layers...)) end @@ -37,7 +38,7 @@ struct NeuroTreeConfig <: Architecture actA::Symbol depth::Int ntrees::Int - proj_size::Int + k::Int hidden_size::Int stack_size::Int scaler::Bool @@ -51,7 +52,7 @@ function NeuroTreeConfig(; kwargs...) :actA => :identity, :depth => 4, :ntrees => 32, - :proj_size => 1, + :k => 1, :hidden_size => 1, :stack_size => 1, :scaler => true, @@ -77,7 +78,7 @@ function NeuroTreeConfig(; kwargs...) Symbol(args[:actA]), args[:depth], args[:ntrees], - args[:proj_size], + args[:k], args[:hidden_size], args[:stack_size], args[:scaler], @@ -91,6 +92,7 @@ function _tree_kwargs(config::NeuroTreeConfig) config.tree_type, config.depth, trees=config.ntrees, + # k=config.k, actA=act_dict[config.actA], config.scaler, config.init_scale, @@ -104,17 +106,15 @@ function (config::NeuroTreeConfig)(; nfeats, outsize, kwargs...) iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") head_outsize = outsize ÷ 2 chain = Chain( - BatchNorm(nfeats, track_stats=false), Parallel( vcat, - StackedNeuroTree(; feats=nfeats, outs=head_outsize, config.hidden_size, config.stack_size, kwargs...), - StackedNeuroTree(; feats=nfeats, outs=head_outsize, config.hidden_size, config.stack_size, kwargs...), + StackedNeuroTree(nfeats => head_outsize; config.hidden_size, config.stack_size, config.k, kwargs...), + StackedNeuroTree(nfeats => head_outsize; config.hidden_size, config.stack_size, config.k, kwargs...), ), ) else chain = Chain( - BatchNorm(nfeats), - StackedNeuroTree(; feats=nfeats, outs=outsize, config.hidden_size, config.stack_size, kwargs...), + StackedNeuroTree(nfeats => outsize; config.hidden_size, config.stack_size, config.k, kwargs...), ) end diff --git a/src/models/embeddings/batchnorm.jl b/src/models/embeddings/batchnorm.jl new file mode 100644 index 0000000..fddf521 --- /dev/null +++ b/src/models/embeddings/batchnorm.jl @@ -0,0 +1,26 @@ +using Lux +using Random +using NNlib + +""" + BatchNormEmbeddings(n_features) + +Embeds each continuous feature via a learned affine transformation followed by +an activation: `activation(w_j * x_j + b_j)`. +Produces a `(d_embedding, n_features, batch)` tensor. + +# Arguments +- `n_features::Int`: Number of input features. +""" +struct BatchNormEmbeddings{L} <: Lux.AbstractLuxWrapperLayer{:layer} + layer::L +end + +function BatchNormEmbeddings(n_features::Int) + return BatchNormEmbeddings(BatchNorm(n_features)) +end + +function (l::BatchNormEmbeddings)(x::AbstractMatrix, ps, st) + x_bn, st = l.layer(x, ps, st) + return x_bn, st +end \ No newline at end of file diff --git a/src/models/embeddings/compute_bins.jl b/src/models/embeddings/compute_bins.jl index 1643379..d0afdfd 100644 --- a/src/models/embeddings/compute_bins.jl +++ b/src/models/embeddings/compute_bins.jl @@ -1,5 +1,5 @@ """ - compute_bins(X; n_bins=48) + compute_bins(X; bins=48) Compute quantile-based bin boundaries for `PiecewiseLinearEncoding`/`PiecewiseLinearEmbeddings`. @@ -8,28 +8,28 @@ model input convention `(n_features, batch)`. Transpose your data before calling # Arguments - `X::AbstractMatrix`: Training data of shape `(n_samples, n_features)`. -- `n_bins::Union{Int, Vector{Int}}`: Number of bins per feature (default `48`). +- `bins::Union{Int, Vector{Int}}`: Number of bins per feature (default `48`). A single `Int` applies the same count to all features. A `Vector{Int}` of length `n_features` specifies per-feature bin counts. # Returns - `Vector{Vector{Float32}}`: Bin edges for each feature. Each vector has between 2 and - `n_bins + 1` elements (fewer if quantiles coincide for low-cardinality features). + `bins + 1` elements (fewer if quantiles coincide for low-cardinality features). """ -function compute_bins(X::AbstractMatrix; n_bins::Union{Int, Vector{Int}}=48) +function compute_bins(X::AbstractMatrix; bins::Union{Int,Vector{Int}}=48) n_samples, n_features = size(X) - n_bins_vec = if n_bins isa Int - @assert n_bins > 1 "n_bins must be > 1, got $n_bins" - @assert n_bins < n_samples "n_bins must be < n_samples, got n_bins=$n_bins, n_samples=$n_samples" - fill(n_bins, n_features) + n_bins_vec = if bins isa Int + @assert bins > 1 "bins must be > 1, got $bins" + @assert bins < n_samples "bins must be < n_samples, got bins=$bins, n_samples=$n_samples" + fill(bins, n_features) else - @assert length(n_bins) == n_features "Length of n_bins must match n_features ($n_features), got $(length(n_bins))" - for (j, nb) in enumerate(n_bins) - @assert nb > 1 "n_bins[$j] must be > 1, got $nb" - @assert nb < n_samples "n_bins[$j] must be < n_samples, got n_bins=$nb, n_samples=$n_samples" + @assert length(bins) == n_features "Length of bins must match n_features ($n_features), got $(length(bins))" + for (j, nb) in enumerate(bins) + @assert nb > 1 "bins[$j] must be > 1, got $nb" + @assert nb < n_samples "bins[$j] must be < n_samples, got bins=$nb, n_samples=$n_samples" end - n_bins + bins end bins = Vector{Vector{Float32}}(undef, n_features) diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl index a5b8bb3..7ba9009 100644 --- a/src/models/embeddings/config.jl +++ b/src/models/embeddings/config.jl @@ -11,9 +11,9 @@ Constructed by the user and passed as `embedding_config` to `NeuroTabRegressor`/ - `d_embedding::Int`: Embedding dimension per feature (default `24`). - `activation`: Activation function after projection (default `relu` for periodic/linear, `identity` for piecewise). -- `n_bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). -- `n_frequencies::Int`: Frequency components for periodic embeddings (default `48`). -- `frequency_init_scale::Float32`: σ for periodic frequency init (default `0.01f0`). +- `bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). +- `frequencies::Int`: Frequency components for periodic embeddings (default `48`). +- `frequencies_init_scale::Float32`: σ for periodic frequencies init (default `0.01f0`). # Callable config(; nfeats, X_train=nothing)-Lux.Chain @@ -24,45 +24,54 @@ struct EmbeddingConfig{F} embedding_type::Symbol d_embedding::Int activation::F - n_bins::Union{Int, Vector{Int}} - n_frequencies::Int - frequency_init_scale::Float32 + bins::Union{Int,Vector{Int}} + frequencies::Int + frequencies_init_scale::Float32 end function EmbeddingConfig(; - embedding_type::Symbol=:periodic, - d_embedding::Int=24, + embedding_type=:periodic, + d_embedding::Int=16, activation=nothing, - n_bins::Union{Int, Vector{Int}}=48, - n_frequencies::Int=48, - frequency_init_scale::Float32=0.01f0, + bins::Union{Int,Vector{Int}}=32, + frequencies::Int=32, + frequencies_init_scale::Float32=0.01f0, ) + + embedding_type = Symbol(embedding_type) # Default activation depends on embedding type if isnothing(activation) activation = embedding_type == :piecewise ? identity : relu end + # Override d_embedding for batchnorm to 1 (for proper derivation of the number of input feature to core model block) + if embedding_type == :batchnorm + d_embedding = 1 + end + return EmbeddingConfig( embedding_type, d_embedding, activation, - n_bins, n_frequencies, frequency_init_scale, + bins, frequencies, frequencies_init_scale, ) end -function (config::EmbeddingConfig)(; nfeats::Int, X_train=nothing) +function (config::EmbeddingConfig)(; nfeats::Int, x_train=nothing) emb = if config.embedding_type == :periodic PeriodicEmbeddings(nfeats, config.d_embedding; - n_frequencies=config.n_frequencies, - frequency_init_scale=config.frequency_init_scale, + frequencies=config.frequencies, + frequencies_init_scale=config.frequencies_init_scale, activation=config.activation) elseif config.embedding_type == :linear LinearEmbeddings(nfeats, config.d_embedding; activation=config.activation) elseif config.embedding_type == :piecewise - @assert X_train !== nothing "Piecewise embeddings require `X_train` to compute bin edges." - bins = compute_bins(X_train; n_bins=config.n_bins) + @assert x_train !== nothing "Piecewise embeddings require `x_train` to compute bin edges." + bins = compute_bins(x_train; bins=config.bins) @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))" PiecewiseLinearEmbeddings(bins, config.d_embedding; activation=config.activation) + elseif config.embedding_type == :batchnorm + BatchNormEmbeddings(nfeats) else error("Unsupported embedding type: $(config.embedding_type)") end return Chain(emb, FlattenLayer()) -end \ No newline at end of file +end diff --git a/src/models/embeddings/embeddings.jl b/src/models/embeddings/embeddings.jl index b24f672..f97eb40 100644 --- a/src/models/embeddings/embeddings.jl +++ b/src/models/embeddings/embeddings.jl @@ -18,6 +18,7 @@ include("nlinear.jl") include("linear.jl") include("periodic.jl") include("piecewise_linear.jl") +include("batchnorm.jl") include("config.jl") end \ No newline at end of file diff --git a/src/models/embeddings/linear.jl b/src/models/embeddings/linear.jl index 65af699..6c2bc65 100644 --- a/src/models/embeddings/linear.jl +++ b/src/models/embeddings/linear.jl @@ -3,34 +3,34 @@ using Random using NNlib """ - LinearEmbeddings(n_features, d_embedding; activation=relu) + LinearEmbeddings(nfeats, d_embedding; activation=relu) Embeds each continuous feature via a learned affine transformation followed by an activation: `activation(w_j * x_j + b_j)`. -Produces a `(d_embedding, n_features, batch)` tensor. +Produces a `(d_embedding, nfeats, batch)` tensor. # Arguments -- `n_features::Int`: Number of input features. +- `nfeats::Int`: Number of input features. - `d_embedding::Int`: Embedding dimension per feature. - `activation`: Activation function applied element-wise (default `relu`). E.g. `relu`, `tanh`, `identity`. """ struct LinearEmbeddings{F} <: Lux.AbstractLuxLayer - n_features::Int + nfeats::Int d_embedding::Int activation::F end -function LinearEmbeddings(n_features::Int, d_embedding::Int; activation=NNlib.relu) - return LinearEmbeddings(n_features, d_embedding, activation) +function LinearEmbeddings(nfeats::Int, d_embedding::Int; activation=NNlib.relu) + return LinearEmbeddings(nfeats, d_embedding, activation) end function Lux.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) limit = Float32(l.d_embedding)^(-0.5f0) - weight = reshape((rand(rng, Float32, l.d_embedding, l.n_features) .* 2f0 .* limit) .- limit, - l.d_embedding, l.n_features, 1) - bias = reshape((rand(rng, Float32, l.d_embedding, l.n_features) .* 2f0 .* limit) .- limit, - l.d_embedding, l.n_features, 1) + weight = reshape((rand(rng, Float32, l.d_embedding, l.nfeats) .* 2f0 .* limit) .- limit, + l.d_embedding, l.nfeats, 1) + bias = reshape((rand(rng, Float32, l.d_embedding, l.nfeats) .* 2f0 .* limit) .- limit, + l.d_embedding, l.nfeats, 1) return (weight=weight, bias=bias) end diff --git a/src/models/embeddings/nlinear.jl b/src/models/embeddings/nlinear.jl index 42f97b8..fc80d5f 100644 --- a/src/models/embeddings/nlinear.jl +++ b/src/models/embeddings/nlinear.jl @@ -38,7 +38,7 @@ end Lux.initialstates(::AbstractRNG, ::NLinear) = (;) -function (l::NLinear)(x::AbstractArray{T, 3}, ps, st) where T +function (l::NLinear)(x::AbstractArray{T,3}, ps, st) where T x_perm = PermutedDimsArray(x, (1, 3, 2)) out = batched_mul(ps.weight, x_perm) diff --git a/src/models/embeddings/periodic.jl b/src/models/embeddings/periodic.jl index 73fbdaa..503f763 100644 --- a/src/models/embeddings/periodic.jl +++ b/src/models/embeddings/periodic.jl @@ -3,26 +3,26 @@ using Random using NNlib """ - Periodic(n_features, n_frequencies, sigma) + Periodic(nfeats, n_frequencies, sigma) Maps each feature to `2 * n_frequencies` sinusoidal components: `[cos(2π w x), sin(2π w x)]`. -Output shape `(2 * n_frequencies, n_features, batch)`. +Output shape `(2 * n_frequencies, nfeats, batch)`. # Arguments -- `n_features::Int`: Number of input features. +- `nfeats::Int`: Number of input features. - `n_frequencies::Int`: Number of frequency components per feature. - `sigma::Float32`: Std-dev for the frequency weight initialization (clamped to ±3σ). """ struct Periodic <: Lux.AbstractLuxLayer - n_features::Int + nfeats::Int n_frequencies::Int sigma::Float32 end function Lux.initialparameters(rng::AbstractRNG, l::Periodic) bound = l.sigma * 3f0 - w = clamp.(l.sigma .* randn(rng, Float32, l.n_frequencies, l.n_features), -bound, bound) - w = reshape(2f0 * Float32(π) .* w, l.n_frequencies, l.n_features, 1) + w = clamp.(l.sigma .* randn(rng, Float32, l.n_frequencies, l.nfeats), -bound, bound) + w = reshape(2f0 * Float32(π) .* w, l.n_frequencies, l.nfeats, 1) return (weight=w,) end @@ -35,17 +35,17 @@ function (l::Periodic)(x::AbstractMatrix, ps, st) end """ - PeriodicEmbeddings(n_features, d_embedding=24; n_frequencies=48, - frequency_init_scale=0.01f0, activation=relu, lite=false) + PeriodicEmbeddings(nfeats, d_embedding=24; n_frequencies=48, + frequencies_init_scale=0.01f0, activation=relu, lite=false) Periodic sinusoidal encoding followed by a learned linear projection. Applies `Periodic` → `NLinear` (or `Dense` if `lite`) → activation. # Arguments -- `n_features::Int`: Number of input features. +- `nfeats::Int`: Number of input features. - `d_embedding::Int`: Output embedding dimension per feature (default `24`). - `n_frequencies::Int`: Sinusoidal frequency components per feature (default `48`). -- `frequency_init_scale::Float32`: σ for frequency weight init (default `0.01f0`). +- `frequencies_init_scale::Float32`: σ for frequency weight init (default `0.01f0`). - `activation`: Activation function applied after projection (default `relu`). E.g. `relu`, `tanh`, `identity`. - `lite::Bool`: Use a single shared `Dense` instead of per-feature `NLinear` (default `false`). Only valid when `activation` is not `identity`. @@ -58,21 +58,21 @@ struct PeriodicEmbeddings{P,L,F} <: Lux.AbstractLuxContainerLayer{(:periodic, :l end function PeriodicEmbeddings( - n_features::Int, + nfeats::Int, d_embedding::Int=24; - n_frequencies::Int=48, - frequency_init_scale::Float32=0.01f0, + frequencies::Int=48, + frequencies_init_scale::Float32=0.01f0, activation=NNlib.relu, lite::Bool=false, ) if lite && activation === identity error("lite=true requires a non-identity activation function") end - periodic = Periodic(n_features, n_frequencies, frequency_init_scale) + periodic = Periodic(nfeats, frequencies, frequencies_init_scale) linear = if lite - Dense(2 * n_frequencies => d_embedding) + Dense(2 * frequencies => d_embedding) else - NLinear(n_features, 2 * n_frequencies, d_embedding) + NLinear(nfeats, 2 * frequencies, d_embedding) end return PeriodicEmbeddings(periodic, linear, activation, lite) end diff --git a/src/models/embeddings/piecewise_linear.jl b/src/models/embeddings/piecewise_linear.jl index b20cdaf..a051cc9 100644 --- a/src/models/embeddings/piecewise_linear.jl +++ b/src/models/embeddings/piecewise_linear.jl @@ -6,33 +6,33 @@ using NNlib PiecewiseLinearEncoding(bins) Non-trainable piecewise-linear encoding using precomputed bin edges. -Output shape `(max_n_bins, n_features, batch)`. +Output shape `(max_n_bins, nfeats, batch)`. # Arguments - `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). """ struct PiecewiseLinearEncoding <: Lux.AbstractLuxLayer bins::Vector{Vector{Float32}} - n_features::Int + nfeats::Int max_n_bins::Int end function PiecewiseLinearEncoding(bins::Vector{<:AbstractVector}) @assert length(bins) > 0 - n_features = length(bins) + nfeats = length(bins) n_bins_list = [length(b) - 1 for b in bins] max_n_bins = maximum(n_bins_list) bins_f32 = [Float32.(b) for b in bins] - return PiecewiseLinearEncoding(bins_f32, n_features, max_n_bins) + return PiecewiseLinearEncoding(bins_f32, nfeats, max_n_bins) end Lux.initialparameters(::AbstractRNG, ::PiecewiseLinearEncoding) = (;) function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) - M, N = l.max_n_bins, l.n_features + M, N = l.max_n_bins, l.nfeats weight = zeros(Float32, M, N) - bias = zeros(Float32, M, N) + bias = zeros(Float32, M, N) for (i, bin_edges) in enumerate(l.bins) bin_width = diff(bin_edges) @@ -44,17 +44,17 @@ function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) # remaining bins fill rows 1:nb-1. Unused rows stay zero # and are clamped to [0, 1] harmlessly. weight[end, i] = w[end] - bias[end, i] = b[end] + bias[end, i] = b[end] if nb > 1 weight[1:nb-1, i] = w[1:end-1] - bias[1:nb-1, i] = b[1:end-1] + bias[1:nb-1, i] = b[1:end-1] end end # Pre-reshape to 3D to avoid per-call reshape in forward pass return ( - weight = reshape(weight, M, N, 1), - bias = reshape(bias, M, N, 1), + weight=reshape(weight, M, N, 1), + bias=reshape(bias, M, N, 1), ) end @@ -71,7 +71,7 @@ end Learnable embeddings on top of `PiecewiseLinearEncoding`. Version `:A`: PLE -> NLinear (with bias). Version `:B`: PLE -> NLinear (zero-init, no bias) + LinearEmbeddings residual. -Output shape `(d_embedding, n_features, batch)`. +Output shape `(d_embedding, nfeats, batch)`. # Arguments - `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). @@ -94,13 +94,13 @@ function PiecewiseLinearEmbeddings( version::Symbol=:B, ) @assert version in (:A, :B) - n_features = length(bins) + nfeats = length(bins) max_n_bins = maximum(length(b) - 1 for b in bins) encoding = PiecewiseLinearEncoding(bins) # Residual path uses raw affine output (no activation) - linear0 = (version == :B) ? LinearEmbeddings(n_features, d_embedding; activation=identity) : nothing - linear = NLinear(n_features, max_n_bins, d_embedding; bias=(version == :A)) + linear0 = (version == :B) ? LinearEmbeddings(nfeats, d_embedding; activation=identity) : nothing + linear = NLinear(nfeats, max_n_bins, d_embedding; bias=(version == :A)) return PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version) end diff --git a/src/models/models.jl b/src/models/models.jl index 84ccf56..ad8b5cd 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -1,8 +1,8 @@ module Models export NeuroTabModel, Architecture -export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig export Embeddings +export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig, MOETreeConfig using ..Losses using Lux: Chain @@ -35,6 +35,9 @@ using .Embeddings include("NeuroTree/neurotrees.jl") using .NeuroTrees +include("MOETree/moetrees.jl") +using .MOETrees + include("TabM/TabM.jl") using .TabM From a29a4e954a2d54e4c5a9c7f917ed4a6d0243dbaf Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Wed, 15 Apr 2026 19:21:29 -0400 Subject: [PATCH 088/120] up --- src/models/embeddings/config.jl | 5 +--- test/core.jl | 45 ++++++++++++++++++++------------- test/embedding.jl | 34 ++++++++++++------------- 3 files changed, 46 insertions(+), 38 deletions(-) diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl index 7ba9009..84902a5 100644 --- a/src/models/embeddings/config.jl +++ b/src/models/embeddings/config.jl @@ -49,10 +49,7 @@ function EmbeddingConfig(; d_embedding = 1 end - return EmbeddingConfig( - embedding_type, d_embedding, activation, - bins, frequencies, frequencies_init_scale, - ) + return EmbeddingConfig(embedding_type, d_embedding, activation, bins, frequencies, frequencies_init_scale) end function (config::EmbeddingConfig)(; nfeats::Int, x_train=nothing) diff --git a/test/core.jl b/test/core.jl index 335b83c..eb4c74a 100644 --- a/test/core.jl +++ b/test/core.jl @@ -1,5 +1,4 @@ -@testset "Core - data iterators" begin -end +@testset "Core - data iterators" begin end @testset "Core - internals test" begin @@ -35,13 +34,14 @@ end ) m = NeuroTabModel(L, chain, info) - + end @testset "Regression - NeuroTree" begin Random.seed!(123) - X, y = rand(1000, 10), randn(1000) + X = randn(Float32, 1000, 10) + y = X[:, 1] .+ 0.5f0 .* X[:, 2] .+ 0.1f0 .* randn(Float32, 1000) df = DataFrame(X, :auto) df[!, :y] = y target_name = "y" @@ -55,8 +55,7 @@ end learner = NeuroTabRegressor(; arch_name="NeuroTreeConfig", - arch_config=Dict( - :depth => 3), + arch_config=Dict(:depth => 3), loss=:mse, nrounds=20, early_stopping_rounds=2, @@ -76,11 +75,18 @@ end target_name, feature_names, deval, + print_every_n=5 ) + p = m(deval) + @test size(p, 1) == nrow(deval) + @test !any(isnan, p) + mse_model = mean((p .- deval.y) .^ 2) + mse_baseline = mean((mean(dtrain.y) .- deval.y) .^ 2) + @test mse_model < mse_baseline end -@testset "Regression - TabM $at" for at in [:tabm, :tabm_mini, :tabm_packed] +@testset "Regression - TabM $arch_type" for arch_type in [:tabm, :tabm_mini, :tabm_packed] Random.seed!(123) X = randn(Float32, 1000, 10) @@ -90,11 +96,13 @@ end target_name = "y" feature_names = setdiff(names(df), [target_name]) - train_indices = 1:800 + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + dtrain = df[train_indices, :] - deval = df[801:end, :] + deval = df[setdiff(1:nrow(df), train_indices), :] - arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=2, d_block=32, dropout=0.0, arch_type=at) + arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=2, d_block=32, dropout=0.0, arch_type) learner = NeuroTabRegressor(arch; loss=:mse, nrounds=20, early_stopping_rounds=2, lr=1e-2) m = NeuroTabModels.fit( @@ -110,6 +118,7 @@ end target_name, feature_names, deval, + print_every_n=5 ) p = m(deval) @@ -140,9 +149,9 @@ end arch_name="NeuroTreeConfig", arch_config=Dict( :depth => 4), - nrounds=100, - batchsize=64, - early_stopping_rounds=10, + embedding_config=Dict(:embedding_type => :batchnorm), + nrounds=200, + early_stopping_rounds=5, lr=3e-2, ) @@ -161,7 +170,7 @@ end end -@testset "Classification - TabM $at" for at in [:tabm, :tabm_mini, :tabm_packed] +@testset "Classification - TabM $arch_type" for arch_type in [:tabm, :tabm_mini, :tabm_packed] Random.seed!(123) X, y = @load_crabs @@ -176,11 +185,12 @@ end dtrain = df[train_indices, :] deval = df[setdiff(1:nrow(df), train_indices), :] - arch = NeuroTabModels.TabMConfig(; k=2, n_blocks=2, d_block=16, dropout=0.0, arch_type=at) + arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=2, d_block=16, dropout=0.0, arch_type) learner = NeuroTabClassifier(arch; - nrounds=500, + embedding_config=Dict(:embedding_type => :batchnorm), + nrounds=200, batchsize=32, - early_stopping_rounds=50, + early_stopping_rounds=5, lr=1e-2, ) @@ -190,6 +200,7 @@ end deval, target_name, feature_names, + print_every_n=5 ) ptrain = [argmax(x) for x in eachrow(m(dtrain))] diff --git a/test/embedding.jl b/test/embedding.jl index 06d494f..f02441c 100644 --- a/test/embedding.jl +++ b/test/embedding.jl @@ -21,17 +21,19 @@ ("NeuroTree", NeuroTabModels.NeuroTreeConfig(; depth=3)), ] - @testset "$arch_name - $et" for (arch_name, arch) in architectures, - et in [:periodic, :linear, :piecewise] - - embed_kw = Dict(:embedding_type => et, :d_embedding => 8) - if et == :piecewise - embed_kw[:n_bins] = 16 + @testset "$arch_name - $embedding_type" for (arch_name, arch) in architectures, + embedding_type in [:periodic, :linear, :piecewise] + + embedding_config = Dict(:embedding_type => embedding_type, :d_embedding => 8) + if embedding_type == :piecewise + embedding_config[:bins] = 16 + elseif embedding_type == :periodic + embedding_config[:frequencies] = 8 end learner = NeuroTabRegressor(arch; loss=:mse, nrounds=20, lr=1e-2, - embedding_config=embed_kw) + embedding_config) m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names) @@ -57,21 +59,19 @@ end dtrain = df[train_indices, :] deval = df[setdiff(1:nrow(df), train_indices), :] - @testset "$et" for et in [:periodic, :linear, :piecewise] - - - embed_kw = Dict(:embedding_type => et, :d_embedding => 8) - if et == :piecewise - embed_kw[:n_bins] = 16 - end - if et == :periodic - embed_kw[:n_frequencies] = 8 + @testset "$embedding_type" for embedding_type in [:periodic, :linear, :piecewise] + + embedding_config = Dict(:embedding_type => embedding_type, :d_embedding => 8) + if embedding_type == :piecewise + embedding_config[:bins] = 16 + elseif embedding_type == :periodic + embedding_config[:frequencies] = 8 end arch = NeuroTabModels.TabMConfig(; k=2, n_blocks=2, d_block=16, dropout=0.0, scaling_init=:random_signs) learner = NeuroTabClassifier(arch; nrounds=500, batchsize=32, early_stopping_rounds=100, lr=1e-2, - embedding_config=embed_kw) + embedding_config) m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names) From 2afe9ef3a5357c153914cb1ca2f18c11e4618e96 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Fri, 17 Apr 2026 11:54:36 -0400 Subject: [PATCH 089/120] up --- benchmarks/YEAR-regression-grp.jl | 131 ++++++++++++++++++++++++++++++ src/Fit/fit.jl | 7 +- src/infer.jl | 23 ++++-- src/learners.jl | 3 + 4 files changed, 152 insertions(+), 12 deletions(-) create mode 100644 benchmarks/YEAR-regression-grp.jl diff --git a/benchmarks/YEAR-regression-grp.jl b/benchmarks/YEAR-regression-grp.jl new file mode 100644 index 0000000..f0075c9 --- /dev/null +++ b/benchmarks/YEAR-regression-grp.jl @@ -0,0 +1,131 @@ +using Random +using CSV +using DataFrames +using Statistics: mean, std +using StatsBase: tiedrank +using NeuroTabModels +using AWS: AWSCredentials, AWSConfig, @service +@service S3 + +aws_creds = AWSCredentials(ENV["AWS_ACCESS_KEY_ID_JDB"], ENV["AWS_SECRET_ACCESS_KEY_JDB"]) +aws_config = AWSConfig(; creds=aws_creds, region="ca-central-1") + +path = "share/data/year/year.csv" +raw = S3.get_object("jeremiedb", path, Dict("response-content-type" => "application/octet-stream"); aws_config) +df = DataFrame(CSV.File(raw, header=false)) +df_tot = copy(df) + +path = "share/data/year/year-train-idx.txt" +raw = S3.get_object("jeremiedb", path, Dict("response-content-type" => "application/octet-stream"); aws_config) +train_idx = DataFrame(CSV.File(raw, header=false))[:, 1] .+ 1 + +path = "share/data/year/year-eval-idx.txt" +raw = S3.get_object("jeremiedb", path, Dict("response-content-type" => "application/octet-stream"); aws_config) +eval_idx = DataFrame(CSV.File(raw, header=false))[:, 1] .+ 1 + +target_name = "y" +rename!(df_tot, "Column1" => target_name) +feature_names = setdiff(names(df_tot), ["y", "w"]) +df_tot.w .= 1.0 + +df_tot.grp = rand(1:round(Int, nrow(df_tot) / 800), nrow(df_tot)) + +# function percent_rank(x::AbstractVector{T}) where {T} +# return tiedrank(x) / (length(x) + 1) +# end +# transform!(df_tot, feature_names .=> percent_rank .=> feature_names) + +dtrain = df_tot[train_idx, :]; +deval = df_tot[eval_idx, :]; +dtest = df_tot[(end-51630+1):end, :]; + + +sort!(dtrain, :grp) +sort!(deval, :grp) +sort!(dtest, :grp) + +arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + actA=:identity, + k=1, + ntrees=32, + depth=4, + stack_size=1, + hidden_size=16, + init_scale=0.1, + scaler=true, +) + +# arch = NeuroTabModels.MOETreeConfig(; +# tree_type=:binary, +# depth=5, +# ntrees=8, +# stack_size=1, +# init_scale=0.1, +# ) + +# arch = NeuroTabModels.TabMConfig(; +# arch_type=:tabm, +# k=16, +# d_block=64, +# n_blocks=3, +# dropout=0.1, +# # scaling_init=:normal, +# ) +# arch = NeuroTabModels.MLPConfig(; +# act=:relu, +# stack_size=1, +# hidden_size=256, +# ) +# arch = NeuroTabModels.ResNetConfig(; +# num_blocks=1, +# hidden_size=128, +# act=:relu, +# dropout=0.5, +# MLE_tree_split=false +# ) + +device = :gpu +loss = :mse # :mse :gaussian_mle :tweedie + +# embedding_config = Dict( +# :embedding_type => :piecewise, +# :d_embedding => 8, +# :activation => nothing, +# :bins => 16, +# :frequencies => 16, +# ) +embedding_config = Dict(:embedding_type => :batchnorm) + +learner = NeuroTabRegressor( + arch; + embedding_config, + loss, + nrounds=200, + early_stopping_rounds=2, + lr=1e-3, + batchsize=512, + device +) + +group_key = "grp" #"grp" # nothing + +@time m = NeuroTabModels.fit( + learner, + dtrain; + deval, + target_name, + feature_names, + group_key, + print_every_n=5, +); + +p_eval = m(deval; device=:cpu); +p_eval = p_eval[:, 1] +mse_eval = mean((p_eval .- deval.y) .^ 2) +@info "MSE - deval" mse_eval + +p_test = m(dtest; device=:cpu); +p_test = p_test[:, 1] +mse_test = mean((p_test .- dtest.y) .^ 2) +@info "MSE - dtest" mse_test diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 63cb023..1f3bfbb 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -51,7 +51,6 @@ function init( lux_loss = get_loss_fn(L) outsize = 1 - scalers = nothing target_levels = nothing target_isordered = false @@ -62,8 +61,10 @@ function init( outsize = length(target_levels) elseif L <: GaussianMLE outsize = 2 - scalers = (mu=mean(df[!, target_name]), sigma=std(df[!, target_name])) - elseif L <: Union{MSE,MAE} + end + + scalers = nothing + if hasproperty(config, :scale_target) && config.scale_target && L <: Union{MSE,MAE,GaussianMLE} scalers = (mu=mean(df[!, target_name]), sigma=std(df[!, target_name])) end diff --git a/src/infer.jl b/src/infer.jl index 9b39654..a73a553 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -35,16 +35,20 @@ _assemble(::Type{<:GaussianMLE}, raw_preds) = reduce(hcat, raw_preds) _assemble(::Type, raw_preds) = vcat([vec(p) for p in raw_preds]...) # Apply inverse link to convert from model scale to natural scale -_inverse_link(::Type{<:LogLoss}, pred, scalers) = sigmoid.(pred) -_inverse_link(::Type{<:Tweedie}, pred, scalers) = exp.(pred) -_inverse_link(::Type{<:Union{MSE,MAE}}, pred, scalers) = pred .* scalers[:sigma] .+ scalers[:mu] - -function _inverse_link(::Type{<:MLogLoss}, pred, scalers) - return Matrix(softmax(pred; dims=1)') +_inverse_link(::Type{<:LogLoss}, pred) = sigmoid.(pred) +_inverse_link(::Type{<:Tweedie}, pred) = exp.(pred) +_inverse_link(::Type{<:Union{MSE,MAE}}, pred) = pred +_inverse_link(::Type{<:MLogLoss}, pred) = Matrix(softmax(pred; dims=1)') +function _inverse_link(::Type{<:GaussianMLE}, pred) + p = Matrix(pred') + @views p[:, 1] .= p[:, 1] + @views p[:, 2] .= exp.(p[:, 2]) + return p end -function _inverse_link(::Type{<:GaussianMLE}, pred, scalers) - p = Matrix(pred') +_scaler(::Type{<:LossType}, p, scalers) = p +_scaler(::Type{<:Union{MSE,MAE}}, p, scalers::NamedTuple) = p .* scalers[:sigma] .+ scalers[:mu] +function _scaler(::Type{<:GaussianMLE}, p, scalers::NamedTuple) @views p[:, 1] .= p[:, 1] .* scalers[:sigma] .+ scalers[:mu] @views p[:, 2] .= exp.(p[:, 2]) .* scalers[:sigma] return p @@ -72,7 +76,8 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L p_raw = _assemble(L, preds) proj || return p_raw - return _inverse_link(L, p_raw, scalers) + p = _inverse_link(L, p_raw) + return _scaler(L, p, scalers) end function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} diff --git a/src/learners.jl b/src/learners.jl index 7bc72d5..37c254f 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -19,6 +19,7 @@ mutable struct NeuroTabRegressor <: MMI.Deterministic wd::Float32 batchsize::Int seed::Int + scale_target::Bool device::Symbol gpuID::Int end @@ -147,6 +148,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) :device => :cpu, :gpuID => 0, :embedding_config => nothing, + :scale_target => true ) args_ignored = setdiff(keys(kwargs), keys(args)) @@ -196,6 +198,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) Float32(args[:wd]), args[:batchsize], args[:seed], + args[:scale_target], device, args[:gpuID] ) From 891d9006851954d3f7dcc38dd28e0c9c8c4c80b4 Mon Sep 17 00:00:00 2001 From: Jeremie Desgagne-Bouchard Date: Fri, 17 Apr 2026 22:35:18 -0400 Subject: [PATCH 090/120] up (#42) --- Project.toml | 2 +- src/Fit/fit.jl | 15 ++++++++++----- src/infer.jl | 16 +++++++++++----- 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/Project.toml b/Project.toml index 4b53e94..46f0759 100644 --- a/Project.toml +++ b/Project.toml @@ -34,7 +34,7 @@ NNlib = "0.9" OneHotArrays = "0.2" Optimisers = "0.4" Random = "1" -Reactant = "=0.2.247" +Reactant = "0.2.247" Statistics = "1" StatsBase = "0.34" Tables = "1.9" diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 1f3bfbb..f885e7c 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -15,7 +15,7 @@ import MLJModelInterface: fit import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam using Lux using Reactant -using Lux: cpu_device, reactant_device +using Lux: cpu_device, gpu_device, reactant_device using DataFrames using CategoricalArrays @@ -23,9 +23,14 @@ using CategoricalArrays include("callback.jl") using .CallBacks -function _get_device(config) - backend = config.device == :gpu ? "gpu" : "cpu" - Reactant.set_default_backend(backend) +""" + _get_device(device::Symbol) + +!!! warning + Returns the default reactant device. + Future behavior should support :cpu/gpu device in addition to the assumed :reactant device. +""" +function _get_device(device::Symbol) return reactant_device() end @@ -44,7 +49,7 @@ function init( offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) group_key = isnothing(group_key) ? nothing : Symbol(group_key) - dev = _get_device(config) + dev = _get_device(config.device) batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) diff --git a/src/infer.jl b/src/infer.jl index a73a553..7e46dcf 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -5,7 +5,7 @@ using ..Losses using ..Models using Lux -using Lux: cpu_device, reactant_device +using Lux: cpu_device, gpu_device, reactant_device using Reactant using Reactant: @compile using NNlib: sigmoid, softmax @@ -18,9 +18,15 @@ export infer, reduce_pred reduce_pred(pred::AbstractMatrix) = pred reduce_pred(pred::AbstractArray{T,3}) where {T} = dropdims(mean(pred; dims=2); dims=2) + +""" + _get_device(device::Symbol) + +!!! warning + Returns the default reactant device. + Future behavior should support :cpu/gpu device in addition to the assumed :reactant device. +""" function _get_device(device::Symbol) - backend = device == :gpu ? "gpu" : "cpu" - Reactant.set_default_backend(backend) return reactant_device() end @@ -55,7 +61,7 @@ function _scaler(::Type{<:GaussianMLE}, p, scalers::NamedTuple) end function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} - dev = _get_device(device) + dev = _get_device(Symbol(device)) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) @@ -81,7 +87,7 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L end function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} - dev = _get_device(device) + dev = _get_device(Symbol(device)) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) From 5fdeecf63351d47e306d59b053daf379836c6b95 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Sun, 19 Apr 2026 02:23:37 -0400 Subject: [PATCH 091/120] test --- src/infer.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/infer.jl b/src/infer.jl index 7e46dcf..a8daf9b 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -103,7 +103,7 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher preds = Vector{AbstractArray}() for (x, mask) in data pred = compiled(m.chain, dev(x), ps, st) - push!(preds, cdev(pred)[mask]) + push!(preds, cdev(pred)[:, mask]) end p_raw = _assemble(L, preds) From bb5460c27dae31c8572a73521e5ca1ae1197c528 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Sun, 19 Apr 2026 02:38:55 -0400 Subject: [PATCH 092/120] up --- src/infer.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/infer.jl b/src/infer.jl index a8daf9b..cc1939d 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -108,7 +108,8 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher p_raw = _assemble(L, preds) proj || return p_raw - return _inverse_link(L, p_raw, scalers) + p = _inverse_link(L, p_raw) + return _scaler(L, p, scalers) end function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, proj::Bool=true) From c05ea3bbeaef6567a1a041a6bfdd1879d2333889 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Wed, 22 Apr 2026 21:03:51 -0400 Subject: [PATCH 093/120] up --- benchmarks/YEAR-regression-grp.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/YEAR-regression-grp.jl b/benchmarks/YEAR-regression-grp.jl index f0075c9..5e3e34e 100644 --- a/benchmarks/YEAR-regression-grp.jl +++ b/benchmarks/YEAR-regression-grp.jl @@ -86,7 +86,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # ) device = :gpu -loss = :mse # :mse :gaussian_mle :tweedie +loss = :gaussian_mle # :mse :gaussian_mle :tweedie # embedding_config = Dict( # :embedding_type => :piecewise, From 84bb1387e988c6cdccbf07e2b518e9a70118d4da Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Wed, 22 Apr 2026 21:23:27 -0400 Subject: [PATCH 094/120] up --- benchmarks/YEAR-regression-grp.jl | 7 ++++++- benchmarks/YEAR-regression.jl | 5 +++++ src/Fit/fit.jl | 2 +- src/data.jl | 2 -- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/benchmarks/YEAR-regression-grp.jl b/benchmarks/YEAR-regression-grp.jl index 5e3e34e..f44228e 100644 --- a/benchmarks/YEAR-regression-grp.jl +++ b/benchmarks/YEAR-regression-grp.jl @@ -35,6 +35,11 @@ df_tot.grp = rand(1:round(Int, nrow(df_tot) / 800), nrow(df_tot)) # end # transform!(df_tot, feature_names .=> percent_rank .=> feature_names) +function norm(x::AbstractVector{T}) where {T} + return (x .- mean(x)) ./ std(x) +end +transform!(df_tot, feature_names .=> norm .=> feature_names) + dtrain = df_tot[train_idx, :]; deval = df_tot[eval_idx, :]; dtest = df_tot[(end-51630+1):end, :]; @@ -86,7 +91,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # ) device = :gpu -loss = :gaussian_mle # :mse :gaussian_mle :tweedie +loss = :mse # :mse :gaussian_mle :tweedie # embedding_config = Dict( # :embedding_type => :piecewise, diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 094001c..ba2a2be 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -33,6 +33,11 @@ df_tot.w .= 1.0 # end # transform!(df_tot, feature_names .=> percent_rank .=> feature_names) +function norm(x::AbstractVector{T}) where {T} + return (x .- mean(x)) ./ std(x) +end +transform!(df_tot, feature_names .=> norm .=> feature_names) + dtrain = df_tot[train_idx, :]; deval = df_tot[eval_idx, :]; dtest = df_tot[(end-51630+1):end, :]; diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index f885e7c..b925c3f 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -166,7 +166,7 @@ function fit( logger = nothing if !isnothing(deval) - cb = CallBack(config, deval, cache; feature_names, target_name, weight_name, offset_name) + cb = CallBack(config, deval, cache; feature_names, target_name, weight_name, offset_name, group_key) logger = init_logger(config) cb(logger, 0, cache[:train_state]) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] diff --git a/src/data.jl b/src/data.jl index cb6dba2..c2dc1bd 100644 --- a/src/data.jl +++ b/src/data.jl @@ -105,7 +105,6 @@ function get_df_loader_train( n = length(dfg) nfeats = length(feature_names) bs = maximum(dfg.ends .- dfg.starts) + 1 - @info "groupedDF loader" bs x = [zeros(Float32, nfeats, bs) for _ in 1:n] y = [zeros(Float32, bs) for _ in 1:n] @@ -180,7 +179,6 @@ function get_df_loader_infer( n = length(dfg) nfeats = length(feature_names) bs = maximum(dfg.ends .- dfg.starts) + 1 - @info "groupedDF loader" bs x = [zeros(Float32, nfeats, bs) for _ in 1:n] mask = [zeros(Bool, bs) for _ in 1:n] From e52f34b0976b82e13e54eff5af55e70bed10cf2b Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 24 Apr 2026 18:58:23 +0530 Subject: [PATCH 095/120] add support for multi-AD backend through extension Signed-off-by: AdityaPandeyCN --- Project.toml | 18 +++++- ext/NeuroTabModelsEnzymeExt.jl | 11 ++++ ext/NeuroTabModelsReactantExt.jl | 47 ++++++++++++++++ ext/NeuroTabModelsZygoteExt.jl | 11 ++++ src/Fit/callback.jl | 19 ++++--- src/Fit/fit.jl | 37 ++++++------ src/infer.jl | 72 +++++++++++++----------- src/learners.jl | 32 +++++++++-- test/core.jl | 96 ++++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 10 files changed, 278 insertions(+), 68 deletions(-) create mode 100644 ext/NeuroTabModelsEnzymeExt.jl create mode 100644 ext/NeuroTabModelsReactantExt.jl create mode 100644 ext/NeuroTabModelsZygoteExt.jl diff --git a/Project.toml b/Project.toml index 46f0759..46ad5e0 100644 --- a/Project.toml +++ b/Project.toml @@ -6,7 +6,6 @@ authors = ["jeremie.db "] [deps] CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" @@ -16,7 +15,6 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" @@ -38,12 +36,26 @@ Reactant = "0.2.247" Statistics = "1" StatsBase = "0.34" Tables = "1.9" +Zygote = "0.7" julia = "1.10" +[weakdeps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[extensions] +NeuroTabModelsEnzymeExt = "Enzyme" +NeuroTabModelsReactantExt = "Reactant" +NeuroTabModelsZygoteExt = "Zygote" + [extras] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d" MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Test", "MLJTestInterface", "MLJBase"] +test = ["Test", "MLJTestInterface", "MLJBase", "Enzyme", "Reactant", "Zygote"] diff --git a/ext/NeuroTabModelsEnzymeExt.jl b/ext/NeuroTabModelsEnzymeExt.jl new file mode 100644 index 0000000..9ba5b1e --- /dev/null +++ b/ext/NeuroTabModelsEnzymeExt.jl @@ -0,0 +1,11 @@ +module NeuroTabModelsEnzymeExt + +using Enzyme +using Lux: AutoEnzyme +using NeuroTabModels + +import NeuroTabModels.Fit: get_ad_backend + +get_ad_backend(::Val{:enzyme}) = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse)) + +end diff --git a/ext/NeuroTabModelsReactantExt.jl b/ext/NeuroTabModelsReactantExt.jl new file mode 100644 index 0000000..34f5cf0 --- /dev/null +++ b/ext/NeuroTabModelsReactantExt.jl @@ -0,0 +1,47 @@ +module NeuroTabModelsReactantExt + +using Lux: Training, reactant_device +using NeuroTabModels +using Reactant +using Reactant: @compile + +import NeuroTabModels.Infer: _get_device, _infer_loop, _infer_grp_loop +import NeuroTabModels.Fit: _single_train_step! +import NeuroTabModels.Fit.CallBacks: _compile_eval_step + +using NeuroTabModels.Infer: _forward_reduce + +_get_device(::Val{:reactant}; gpuID::Integer=0) = reactant_device() + +function _infer_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st) + compiled = @compile _forward_reduce(chain, dev(x0), ps, st) + + preds = Vector{AbstractArray}() + for x in data + if size(x) == size(x0) + pred = compiled(chain, dev(x), ps, st) + else + pred = Reactant.@jit _forward_reduce(chain, dev(x), ps, st) + end + push!(preds, cdev(pred)) + end + return preds +end + +function _infer_grp_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st) + compiled = @compile _forward_reduce(chain, dev(x0), ps, st) + + preds = Vector{AbstractArray}() + for (x, mask) in data + pred = compiled(chain, dev(x), ps, st) + push!(preds, cdev(pred)[:, mask]) + end + return preds +end + +_single_train_step!(::Val{:reactant}, ad_backend, lux_loss, d, ts) = + Training.single_train_step!(ad_backend, lux_loss, d, ts; return_gradients=Val(false)) + +_compile_eval_step(::Val{:reactant}, step, args...) = @compile step(args...) + +end diff --git a/ext/NeuroTabModelsZygoteExt.jl b/ext/NeuroTabModelsZygoteExt.jl new file mode 100644 index 0000000..b5556f3 --- /dev/null +++ b/ext/NeuroTabModelsZygoteExt.jl @@ -0,0 +1,11 @@ +module NeuroTabModelsZygoteExt + +using Lux: AutoZygote +using NeuroTabModels +using Zygote + +import NeuroTabModels.Fit: get_ad_backend + +get_ad_backend(::Val{:zygote}) = AutoZygote() + +end diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index e16de91..5fe4a5c 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -4,12 +4,11 @@ using DataFrames using Statistics: mean, median using ..Learners: LearnerTypes -using ...Infer: reduce_pred +using ...Infer: reduce_pred, _get_device using ..Data: get_df_loader_train using ..Metrics -using Lux: Training, reactant_device, testmode -using Reactant: @compile, set_default_backend +using Lux: Training, testmode export CallBack, init_logger, update_logger!, agg_logger @@ -34,7 +33,7 @@ function CallBack( offset_name=nothing, group_key=nothing ) - dev = reactant_device() + dev = _get_device(config.device; gpuID=config.gpuID) ts = cache[:train_state] scalers = cache[:scalers] batchsize = config.batchsize @@ -45,33 +44,35 @@ function CallBack( ps, st = ts.parameters, testmode(ts.states) d0 = first(deval) - eval_compiled = _compile_eval_step(ts.model, feval, d0, ps, st) + eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.device == :reactant) return CallBack(deval, eval_compiled) end -function _compile_eval_step(chain, feval, d0, ps, st) +function _build_eval_step(chain, feval, d0, ps, st; reactant::Bool) if length(d0) == 2 function _step2(x, y, ps, st) m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y; agg=sum), eltype(y)(last(size(y))) end - return @compile _step2(d0[1], d0[2], ps, st) + return reactant ? _compile_eval_step(Val(:reactant), _step2, d0[1], d0[2], ps, st) : _step2 elseif length(d0) == 3 function _step3(x, y, w, ps, st) m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y, w; agg=sum), sum(w) end - return @compile _step3(d0[1], d0[2], d0[3], ps, st) + return reactant ? _compile_eval_step(Val(:reactant), _step3, d0[1], d0[2], d0[3], ps, st) : _step3 else function _step4(x, y, w, offset, ps, st) m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y, w, offset; agg=sum), sum(w) end - return @compile _step4(d0[1], d0[2], d0[3], d0[4], ps, st) + return reactant ? _compile_eval_step(Val(:reactant), _step4, d0[1], d0[2], d0[3], d0[4], ps, st) : _step4 end end +_compile_eval_step(::Val{D}, step, args...) where {D} = step + function init_logger(config::LearnerTypes) logger = Dict( :name => String(config.metric), diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index b925c3f..953b235 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -7,33 +7,27 @@ using ..Learners using ..Models using ..Losses using ..Metrics -using ..Infer: reduce_pred +using ..Infer: reduce_pred, _get_device import Random: Xoshiro import Statistics: mean, std import MLJModelInterface: fit import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam using Lux -using Reactant -using Lux: cpu_device, gpu_device, reactant_device +using Lux: cpu_device using DataFrames using CategoricalArrays +get_ad_backend(backend::Symbol) = get_ad_backend(Val(backend)) +get_ad_backend(::Val{b}) where {b} = error( + "Unsupported or unloaded `backend=:$b`. Supported: [:enzyme, :zygote]. " * + "`:enzyme` requires `using Enzyme`, `:zygote` requires `using Zygote`." +) + include("callback.jl") using .CallBacks -""" - _get_device(device::Symbol) - -!!! warning - Returns the default reactant device. - Future behavior should support :cpu/gpu device in addition to the assumed :reactant device. -""" -function _get_device(device::Symbol) - return reactant_device() -end - function init( config::LearnerTypes, df::AbstractDataFrame; @@ -49,7 +43,7 @@ function init( offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) group_key = isnothing(group_key) ? nothing : Symbol(group_key) - dev = _get_device(config.device) + dev = _get_device(config.device; gpuID=config.gpuID) batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) @@ -101,7 +95,10 @@ function init( :group_key => group_key, :target_levels => target_levels, :target_isordered => target_isordered, - :scalers => scalers + :scalers => scalers, + :backend => config.backend, + :device => config.device, + :gpuID => config.gpuID ) m = NeuroTabModel(L, chain, info) @@ -110,7 +107,7 @@ function init( opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) - return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts, :scalers => scalers) + return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts, :scalers => scalers, :ad_backend => get_ad_backend(config.backend)) end """ @@ -201,10 +198,14 @@ function _sync_params_to_model!(m, cache) m.info[:st] = cdev(Lux.testmode(ts.states)) end +_single_train_step!(device::Symbol, ad_backend, lux_loss, d, ts) = _single_train_step!(Val(device), ad_backend, lux_loss, d, ts) +_single_train_step!(::Val{D}, ad_backend, lux_loss, d, ts) where {D} = Training.single_train_step!(ad_backend, lux_loss, d, ts) + function fit_iter!(m, cache) ts, lux_loss = cache[:train_state], cache[:lux_loss] + ad_backend = cache[:ad_backend] for d in cache[:data] - _, loss, _, ts = Training.single_train_step!(AutoEnzyme(), lux_loss, d, ts) + _, loss, _, ts = _single_train_step!(m.info[:device], ad_backend, lux_loss, d, ts) end cache[:train_state] = ts m.info[:nrounds] += 1 diff --git a/src/infer.jl b/src/infer.jl index cc1939d..b5996c7 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -5,9 +5,7 @@ using ..Losses using ..Models using Lux -using Lux: cpu_device, gpu_device, reactant_device -using Reactant -using Reactant: @compile +using Lux: cpu_device, gpu_device using NNlib: sigmoid, softmax using Statistics: mean using DataFrames: AbstractDataFrame, GroupedDataFrame, groupby @@ -18,17 +16,19 @@ export infer, reduce_pred reduce_pred(pred::AbstractMatrix) = pred reduce_pred(pred::AbstractArray{T,3}) where {T} = dropdims(mean(pred; dims=2); dims=2) - """ - _get_device(device::Symbol) + _get_device(device::Symbol; gpuID::Integer=0) -!!! warning - Returns the default reactant device. - Future behavior should support :cpu/gpu device in addition to the assumed :reactant device. +Resolve a Lux device from a symbol: `:cpu`, `:gpu`, or `:reactant`. +`:reactant` requires `using Reactant` to activate `NeuroTabModelsReactantExt`. """ -function _get_device(device::Symbol) - return reactant_device() -end +_get_device(device::Symbol; gpuID::Integer=0) = _get_device(Val(device); gpuID) +_get_device(::Val{:cpu}; gpuID::Integer=0) = cpu_device() +_get_device(::Val{:gpu}; gpuID::Integer=0) = gpu_device(gpuID == 0 ? nothing : gpuID) +_get_device(::Val{D}; gpuID::Integer=0) where {D} = error( + "Unsupported `device=:$D`. Supported devices are [:cpu, :gpu, :reactant]. " * + "`:reactant` requires `using Reactant` to activate `NeuroTabModelsReactantExt`." +) function _forward_reduce(chain, x, ps, st) pred, _ = chain(x, ps, st) @@ -60,25 +60,37 @@ function _scaler(::Type{<:GaussianMLE}, p, scalers::NamedTuple) return p end +# Default inference loops for :cpu and :gpu. +# NeuroTabModelsReactantExt adds ::Val{:reactant} methods that run a Reactant-compiled forward pass. +function _infer_loop(::Val, chain, data, x0, dev, cdev, ps, st) + preds = Vector{AbstractArray}() + for x in data + pred = _forward_reduce(chain, dev(x), ps, st) + push!(preds, cdev(pred)) + end + return preds +end + +# Grouped variant: each batch is `(x, mask)`; `mask` selects real rows from padded groups. +function _infer_grp_loop(::Val, chain, data, x0, dev, cdev, ps, st) + preds = Vector{AbstractArray}() + for (x, mask) in data + pred = _forward_reduce(chain, dev(x), ps, st) + push!(preds, cdev(pred)[:, mask]) + end + return preds +end + function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} - dev = _get_device(Symbol(device)) + device = Symbol(device) + dev = _get_device(device) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) scalers = m.info[:scalers] x0 = first(data) - compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) - - preds = Vector{AbstractArray}() - for x in data - if size(x) == size(x0) - pred = compiled(m.chain, dev(x), ps, st) - else - pred = Reactant.@jit _forward_reduce(m.chain, dev(x), ps, st) - end - push!(preds, cdev(pred)) - end + preds = _infer_loop(Val(device), m.chain, data, x0, dev, cdev, ps, st) p_raw = _assemble(L, preds) proj || return p_raw @@ -87,7 +99,8 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L end function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} - dev = _get_device(Symbol(device)) + device = Symbol(device) + dev = _get_device(device) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) @@ -98,13 +111,7 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher # data = data |> dev # (x0, mask0) = first(data) # @info typeof("mask0-dev") mask0 - compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) - - preds = Vector{AbstractArray}() - for (x, mask) in data - pred = compiled(m.chain, dev(x), ps, st) - push!(preds, cdev(pred)[:, mask]) - end + preds = _infer_grp_loop(Val(device), m.chain, data, x0, dev, cdev, ps, st) p_raw = _assemble(L, preds) proj || return p_raw @@ -133,4 +140,5 @@ end # return infer(m, [(x,)]; device) # end -end \ No newline at end of file +end + diff --git a/src/learners.jl b/src/learners.jl index 37c254f..f845efa 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -20,6 +20,7 @@ mutable struct NeuroTabRegressor <: MMI.Deterministic batchsize::Int seed::Int scale_target::Bool + backend::Symbol device::Symbol gpuID::Int end @@ -43,11 +44,14 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl]( - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. -- `device=:cpu`: Device on which to perform the computation, either `:cpu` or `:gpu` -- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `backend=:enzyme`: AD backend used by Lux. One of `:enzyme` or `:zygote`. +- `device=:reactant`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. +- `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`. +`backend=:enzyme` works on `:cpu`, `:gpu`, and `:reactant`; `backend=:zygote` works on `:cpu` and `:gpu`. + # Internal API Do `config = NeuroTabRegressor()` to construct an instance with default hyper-parameters. @@ -145,7 +149,8 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) :wd => 0.0f0, :batchsize => 2048, :seed => 123, - :device => :cpu, + :backend => :enzyme, + :device => :reactant, :gpuID => 0, :embedding_config => nothing, :scale_target => true @@ -179,7 +184,11 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) error("Invalid metric. Must be one of: $_metric_list") end + backend = Symbol(args[:backend]) device = Symbol(args[:device]) + if device == :reactant && backend != :enzyme + error("`device=:reactant` requires `backend=:enzyme`. Got `backend=:$backend`.") + end # Build EmbeddingConfig from Dict if needed embed = args[:embedding_config] @@ -199,6 +208,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) args[:batchsize], args[:seed], args[:scale_target], + backend, device, args[:gpuID] ) @@ -224,6 +234,7 @@ mutable struct NeuroTabClassifier <: MMI.Probabilistic wd::Float32 batchsize::Int seed::Int + backend::Symbol device::Symbol gpuID::Int end @@ -241,10 +252,13 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl] - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. -- `device=:cpu`: Device on which to perform the computation, either `:cpu` or `:gpu` -- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `backend=:enzyme`: AD backend used by Lux. One of `:enzyme` or `:zygote`. +- `device=:reactant`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. +- `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. +`backend=:enzyme` works on `:cpu`, `:gpu`, and `:reactant`; `backend=:zygote` works on `:cpu` and `:gpu`. + # Internal API Do `config = NeuroTabClassifier()` to construct an instance with default hyper-parameters. @@ -341,7 +355,8 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) :wd => 0.0f0, :batchsize => 2048, :seed => 123, - :device => :cpu, + :backend => :enzyme, + :device => :reactant, :gpuID => 0, :embedding_config => nothing, ) @@ -361,7 +376,11 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) args[arg] = kwargs[arg] end + backend = Symbol(args[:backend]) device = Symbol(args[:device]) + if device == :reactant && backend != :enzyme + error("`device=:reactant` requires `backend=:enzyme`. Got `backend=:$backend`.") + end # Build EmbeddingConfig from Dict if needed embed = args[:embedding_config] @@ -380,6 +399,7 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) Float32(args[:wd]), args[:batchsize], args[:seed], + backend, device, args[:gpuID] ) diff --git a/test/core.jl b/test/core.jl index eb4c74a..7605d67 100644 --- a/test/core.jl +++ b/test/core.jl @@ -208,4 +208,100 @@ end @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 @test mean(peval .== levelcode.(deval.class)) > 0.95 +end + +@testset "Backend/device - reactant requires enzyme" begin + @test_throws ErrorException NeuroTabRegressor(; backend=:zygote, device=:reactant) + @test_throws ErrorException NeuroTabClassifier(; backend=:zygote, device=:reactant) +end + +@testset "Backend/device - Regression ($backend, $device)" for (backend, device) in [ + (:enzyme, :cpu), + (:zygote, :cpu), + (:enzyme, :reactant), +] + + Random.seed!(123) + X = randn(Float32, 1000, 10) + y = X[:, 1] .+ 0.5f0 .* X[:, 2] .+ 0.1f0 .* randn(Float32, 1000) + df = DataFrame(X, :auto) + df[!, :y] = y + target_name = "y" + feature_names = setdiff(names(df), [target_name]) + + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + learner = NeuroTabRegressor(; + arch_name="NeuroTreeConfig", + arch_config=Dict(:depth => 3), + loss=:mse, + nrounds=20, + early_stopping_rounds=2, + lr=1e-1, + backend, + device, + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + target_name, + feature_names, + deval, + ) + + p = m(deval) + @test size(p, 1) == nrow(deval) + @test !any(isnan, p) + mse_model = mean((p .- deval.y) .^ 2) + mse_baseline = mean((mean(dtrain.y) .- deval.y) .^ 2) + @test mse_model < mse_baseline +end + +@testset "Backend/device - Classification ($backend, $device)" for (backend, device) in [ + (:enzyme, :cpu), + (:zygote, :cpu), +] + + Random.seed!(123) + X, y = @load_crabs + df = DataFrame(X) + df[!, :class] = y + target_name = "class" + feature_names = setdiff(names(df), [target_name]) + + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + learner = NeuroTabClassifier(; + arch_name="NeuroTreeConfig", + arch_config=Dict(:depth => 4), + embedding_config=Dict(:embedding_type => :batchnorm), + nrounds=200, + early_stopping_rounds=5, + lr=3e-2, + backend, + device, + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + deval, + target_name, + feature_names, + ) + + ptrain = [argmax(x) for x in eachrow(m(dtrain))] + peval = [argmax(x) for x in eachrow(m(deval))] + @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 + @test mean(peval .== levelcode.(deval.class)) > 0.95 + end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 8959b2b..a60810b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,6 +8,9 @@ using StatsBase: sample using Random using MLJBase using MLJTestInterface +using Enzyme +using Reactant +using Zygote include("core.jl") include("embedding.jl") From 8944fd94d8c6a55675955406aea2a06a5303be35 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Fri, 24 Apr 2026 14:21:38 -0400 Subject: [PATCH 096/120] up --- benchmarks/benchmark_mse.jl | 50 +++++++++++++++++++-------------- src/models/embeddings/config.jl | 28 ++++++++++++------ 2 files changed, 48 insertions(+), 30 deletions(-) diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 28b3772..65294b6 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -16,36 +16,44 @@ feature_names = names(dtrain) dtrain.y = Y target_name = "y" -# arch = NeuroTabModels.NeuroTreeConfig(; -# tree_type=:binary, -# actA=:identity, -# init_scale=1.0, -# depth=4, -# ntrees=32, -# stack_size=1, -# hidden_size=1, -# scaler=false, -# ) -arch = NeuroTabModels.TabMConfig(; - arch_type=:tabm, - k=16, - d_block=64, - n_blocks=3, - dropout=0.1, - bins=nothing, - use_embeddings=false, - embedding_type=:periodic, - d_embedding=16, - scaling_init=:random_signs, +arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + actA=:identity, + init_scale=1.0, + depth=4, + ntrees=32, + stack_size=1, + hidden_size=1, + scaler=false, ) +# arch = NeuroTabModels.TabMConfig(; +# arch_type=:tabm, +# k=16, +# d_block=64, +# n_blocks=3, +# dropout=0.1, +# bins=nothing, +# use_embeddings=false, +# embedding_type=:periodic, +# d_embedding=16, +# scaling_init=:random_signs, +# ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, # stack_size=1, # hidden_size=64, # ) +# embedding_config = Dict( +# :embedding_type => :linear, +# :d_embedding => 8, +# :activation => "identity", +# ) +embedding_config = Dict(:embedding_type => :batchnorm) + learner = NeuroTabRegressor( arch; + embedding_config, loss=:mse, nrounds=10, lr=1e-2, diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl index 84902a5..8d902bc 100644 --- a/src/models/embeddings/config.jl +++ b/src/models/embeddings/config.jl @@ -1,4 +1,13 @@ -using NNlib: relu +using NNlib: relu, tanh, softplus, hardtanh, tanhshrink + +const act_dict = Dict( + :identity => identity, + :relu => relu, + :tanh => tanh, + :softplus => softplus, + :hardtanh => hardtanh, + :tanhshrink => tanhshrink, +) """ EmbeddingConfig(; kwargs...) @@ -9,8 +18,8 @@ Constructed by the user and passed as `embedding_config` to `NeuroTabRegressor`/ # Arguments - `embedding_type::Symbol`: `:periodic`, `:linear`, or `:piecewise` (default `:periodic`). - `d_embedding::Int`: Embedding dimension per feature (default `24`). -- `activation`: Activation function after projection (default `relu` for periodic/linear, - `identity` for piecewise). +- `activation::Symbol`: activation function symbol (default `:relu` for periodic/linear, + `:identity` for piecewise). - `bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). - `frequencies::Int`: Frequency components for periodic embeddings (default `48`). - `frequencies_init_scale::Float32`: σ for periodic frequencies init (default `0.01f0`). @@ -20,10 +29,10 @@ Constructed by the user and passed as `embedding_config` to `NeuroTabRegressor`/ Returns a `Chain(embedding_layer, FlattenLayer())` ready to prepend to any backbone. """ -struct EmbeddingConfig{F} +struct EmbeddingConfig embedding_type::Symbol d_embedding::Int - activation::F + activation::Symbol bins::Union{Int,Vector{Int}} frequencies::Int frequencies_init_scale::Float32 @@ -41,8 +50,9 @@ function EmbeddingConfig(; embedding_type = Symbol(embedding_type) # Default activation depends on embedding type if isnothing(activation) - activation = embedding_type == :piecewise ? identity : relu + activation = embedding_type == :piecewise ? :identity : :relu end + activation = Symbol(activation) # Override d_embedding for batchnorm to 1 (for proper derivation of the number of input feature to core model block) if embedding_type == :batchnorm @@ -57,14 +67,14 @@ function (config::EmbeddingConfig)(; nfeats::Int, x_train=nothing) PeriodicEmbeddings(nfeats, config.d_embedding; frequencies=config.frequencies, frequencies_init_scale=config.frequencies_init_scale, - activation=config.activation) + activation=act_dict[config.activation]) elseif config.embedding_type == :linear - LinearEmbeddings(nfeats, config.d_embedding; activation=config.activation) + LinearEmbeddings(nfeats, config.d_embedding; activation=act_dict[config.activation]) elseif config.embedding_type == :piecewise @assert x_train !== nothing "Piecewise embeddings require `x_train` to compute bin edges." bins = compute_bins(x_train; bins=config.bins) @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))" - PiecewiseLinearEmbeddings(bins, config.d_embedding; activation=config.activation) + PiecewiseLinearEmbeddings(bins, config.d_embedding; activation=act_dict[config.activation]) elseif config.embedding_type == :batchnorm BatchNormEmbeddings(nfeats) else From 0d1ff5e4821fe4e9dda7d69d01cd89d85b496d41 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sat, 25 Apr 2026 18:11:55 +0530 Subject: [PATCH 097/120] revert default to cpu Signed-off-by: AdityaPandeyCN --- src/learners.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/learners.jl b/src/learners.jl index f845efa..73f60ca 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -45,7 +45,7 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl]( - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. - `backend=:enzyme`: AD backend used by Lux. One of `:enzyme` or `:zygote`. -- `device=:reactant`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. +- `device=:cpu`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. - `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`. @@ -150,7 +150,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) :batchsize => 2048, :seed => 123, :backend => :enzyme, - :device => :reactant, + :device => :cpu, :gpuID => 0, :embedding_config => nothing, :scale_target => true @@ -253,7 +253,7 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl] - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. - `backend=:enzyme`: AD backend used by Lux. One of `:enzyme` or `:zygote`. -- `device=:reactant`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. +- `device=:cpu`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. - `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. @@ -356,7 +356,7 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) :batchsize => 2048, :seed => 123, :backend => :enzyme, - :device => :reactant, + :device => :cpu, :gpuID => 0, :embedding_config => nothing, ) From 5f60665820023981eda2e0c5770eb9d5d924a441 Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Sat, 25 Apr 2026 18:51:46 +0530 Subject: [PATCH 098/120] REMOVE ZYGOTE COMPAT Signed-off-by: AdityaPandeyCN --- Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/Project.toml b/Project.toml index 46ad5e0..3ee8dbc 100644 --- a/Project.toml +++ b/Project.toml @@ -36,7 +36,6 @@ Reactant = "0.2.247" Statistics = "1" StatsBase = "0.34" Tables = "1.9" -Zygote = "0.7" julia = "1.10" [weakdeps] From 1eb6420d1ba4d40ba9f8993bd57ac81810a7bc5a Mon Sep 17 00:00:00 2001 From: AdityaPandeyCN Date: Fri, 1 May 2026 18:23:27 +0530 Subject: [PATCH 099/120] Move Reactant selection from to , keep Enzyme as the AD backend underneath, and default learners to Zygote/GPU for broader compatibility. Signed-off-by: AdityaPandeyCN --- Project.toml | 3 ++- ext/NeuroTabModelsReactantExt.jl | 5 ++++- src/Fit/callback.jl | 4 ++-- src/Fit/fit.jl | 9 ++++---- src/infer.jl | 33 ++++++++++++++++------------- src/learners.jl | 36 +++++++++++++++++++------------- test/core.jl | 4 ++-- 7 files changed, 55 insertions(+), 39 deletions(-) diff --git a/Project.toml b/Project.toml index 3ee8dbc..61e6f2e 100644 --- a/Project.toml +++ b/Project.toml @@ -32,10 +32,11 @@ NNlib = "0.9" OneHotArrays = "0.2" Optimisers = "0.4" Random = "1" -Reactant = "0.2.247" +Reactant = "0.2.254" Statistics = "1" StatsBase = "0.34" Tables = "1.9" +Zygote = "0.7" julia = "1.10" [weakdeps] diff --git a/ext/NeuroTabModelsReactantExt.jl b/ext/NeuroTabModelsReactantExt.jl index 34f5cf0..791cc74 100644 --- a/ext/NeuroTabModelsReactantExt.jl +++ b/ext/NeuroTabModelsReactantExt.jl @@ -11,7 +11,10 @@ import NeuroTabModels.Fit.CallBacks: _compile_eval_step using NeuroTabModels.Infer: _forward_reduce -_get_device(::Val{:reactant}; gpuID::Integer=0) = reactant_device() +function _get_device(::Val{:reactant}, ::Val{D}; gpuID::Integer=0) where {D} + Reactant.set_default_backend(String(D)) + return reactant_device() +end function _infer_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st) compiled = @compile _forward_reduce(chain, dev(x0), ps, st) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 5fe4a5c..04b7f7b 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -33,7 +33,7 @@ function CallBack( offset_name=nothing, group_key=nothing ) - dev = _get_device(config.device; gpuID=config.gpuID) + dev = _get_device(config.backend, config.device; gpuID=config.gpuID) ts = cache[:train_state] scalers = cache[:scalers] batchsize = config.batchsize @@ -44,7 +44,7 @@ function CallBack( ps, st = ts.parameters, testmode(ts.states) d0 = first(deval) - eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.device == :reactant) + eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.backend == :reactant) return CallBack(deval, eval_compiled) end diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index 953b235..cf9e41b 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -20,9 +20,10 @@ using DataFrames using CategoricalArrays get_ad_backend(backend::Symbol) = get_ad_backend(Val(backend)) +get_ad_backend(::Val{:reactant}) = get_ad_backend(Val(:enzyme)) get_ad_backend(::Val{b}) where {b} = error( - "Unsupported or unloaded `backend=:$b`. Supported: [:enzyme, :zygote]. " * - "`:enzyme` requires `using Enzyme`, `:zygote` requires `using Zygote`." + "Unsupported or unloaded `backend=:$b`. Supported: [:enzyme, :zygote, :reactant]. " * + "`:enzyme` and `:reactant` require `using Enzyme`, `:zygote` requires `using Zygote`." ) include("callback.jl") @@ -43,7 +44,7 @@ function init( offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) group_key = isnothing(group_key) ? nothing : Symbol(group_key) - dev = _get_device(config.device; gpuID=config.gpuID) + dev = _get_device(config.backend, config.device; gpuID=config.gpuID) batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) @@ -205,7 +206,7 @@ function fit_iter!(m, cache) ts, lux_loss = cache[:train_state], cache[:lux_loss] ad_backend = cache[:ad_backend] for d in cache[:data] - _, loss, _, ts = _single_train_step!(m.info[:device], ad_backend, lux_loss, d, ts) + _, loss, _, ts = _single_train_step!(m.info[:backend], ad_backend, lux_loss, d, ts) end cache[:train_state] = ts m.info[:nrounds] += 1 diff --git a/src/infer.jl b/src/infer.jl index b5996c7..8bb0089 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -19,15 +19,16 @@ reduce_pred(pred::AbstractArray{T,3}) where {T} = dropdims(mean(pred; dims=2); d """ _get_device(device::Symbol; gpuID::Integer=0) -Resolve a Lux device from a symbol: `:cpu`, `:gpu`, or `:reactant`. -`:reactant` requires `using Reactant` to activate `NeuroTabModelsReactantExt`. +Resolve a Lux device from a symbol: `:cpu` or `:gpu`. +`backend=:reactant` requires `using Reactant` to activate `NeuroTabModelsReactantExt`. """ _get_device(device::Symbol; gpuID::Integer=0) = _get_device(Val(device); gpuID) +_get_device(backend::Symbol, device::Symbol; gpuID::Integer=0) = _get_device(Val(backend), Val(device); gpuID) +_get_device(::Val, ::Val{D}; gpuID::Integer=0) where {D} = _get_device(Val(D); gpuID) _get_device(::Val{:cpu}; gpuID::Integer=0) = cpu_device() _get_device(::Val{:gpu}; gpuID::Integer=0) = gpu_device(gpuID == 0 ? nothing : gpuID) _get_device(::Val{D}; gpuID::Integer=0) where {D} = error( - "Unsupported `device=:$D`. Supported devices are [:cpu, :gpu, :reactant]. " * - "`:reactant` requires `using Reactant` to activate `NeuroTabModelsReactantExt`." + "Unsupported `device=:$D`. Supported devices are [:cpu, :gpu]." ) function _forward_reduce(chain, x, ps, st) @@ -81,16 +82,17 @@ function _infer_grp_loop(::Val, chain, data, x0, dev, cdev, ps, st) return preds end -function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} +function infer(m::NeuroTabModel{L}, data; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) where {L} device = Symbol(device) - dev = _get_device(device) + backend = Symbol(backend) + dev = _get_device(backend, device; gpuID=get(m.info, :gpuID, 0)) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) scalers = m.info[:scalers] x0 = first(data) - preds = _infer_loop(Val(device), m.chain, data, x0, dev, cdev, ps, st) + preds = _infer_loop(Val(backend), m.chain, data, x0, dev, cdev, ps, st) p_raw = _assemble(L, preds) proj || return p_raw @@ -98,9 +100,10 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L return _scaler(L, p, scalers) end -function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} +function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) where {L} device = Symbol(device) - dev = _get_device(device) + backend = Symbol(backend) + dev = _get_device(backend, device; gpuID=get(m.info, :gpuID, 0)) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) @@ -111,7 +114,7 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher # data = data |> dev # (x0, mask0) = first(data) # @info typeof("mask0-dev") mask0 - preds = _infer_grp_loop(Val(device), m.chain, data, x0, dev, cdev, ps, st) + preds = _infer_grp_loop(Val(backend), m.chain, data, x0, dev, cdev, ps, st) p_raw = _assemble(L, preds) proj || return p_raw @@ -119,21 +122,21 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher return _scaler(L, p, scalers) end -function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, proj::Bool=true) +function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) group_key = m.info[:group_key] if isnothing(group_key) dinfer = get_df_loader_infer(df; feature_names=m.info[:feature_names], batchsize=2048) - p = infer(m, dinfer; device, proj) + p = infer(m, dinfer; device, backend, proj) else dfg = groupby(df, group_key; sort=true) dinfer = get_df_loader_infer(dfg; feature_names=m.info[:feature_names], batchsize=2048) - p = infer_grp(m, dinfer; device, proj) + p = infer_grp(m, dinfer; device, backend, proj) end return p end -function (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, proj::Bool=true) - return infer(m, df; device, proj) +function (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) + return infer(m, df; device, backend, proj) end # function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) diff --git a/src/learners.jl b/src/learners.jl index 73f60ca..3e4d8d9 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -44,13 +44,14 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl]( - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. -- `backend=:enzyme`: AD backend used by Lux. One of `:enzyme` or `:zygote`. -- `device=:cpu`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. +- `backend=:zygote`: Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`. +- `device=:gpu`: Execution device. One of `:cpu` or `:gpu`. - `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`. -`backend=:enzyme` works on `:cpu`, `:gpu`, and `:reactant`; `backend=:zygote` works on `:cpu` and `:gpu`. +`backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD. +`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. # Internal API @@ -149,8 +150,8 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) :wd => 0.0f0, :batchsize => 2048, :seed => 123, - :backend => :enzyme, - :device => :cpu, + :backend => :zygote, + :device => :gpu, :gpuID => 0, :embedding_config => nothing, :scale_target => true @@ -186,8 +187,11 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) backend = Symbol(args[:backend]) device = Symbol(args[:device]) - if device == :reactant && backend != :enzyme - error("`device=:reactant` requires `backend=:enzyme`. Got `backend=:$backend`.") + if device == :reactant + error("Use `backend=:reactant` with `device=:cpu` or `device=:gpu` instead of `device=:reactant`.") + end + if backend == :enzyme && device == :gpu + @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU." end # Build EmbeddingConfig from Dict if needed @@ -252,12 +256,13 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl] - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. -- `backend=:enzyme`: AD backend used by Lux. One of `:enzyme` or `:zygote`. -- `device=:cpu`: Execution device. One of `:cpu`, `:gpu`, or `:reactant`. +- `backend=:zygote`: Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`. +- `device=:gpu`: Execution device. One of `:cpu` or `:gpu`. - `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. -`backend=:enzyme` works on `:cpu`, `:gpu`, and `:reactant`; `backend=:zygote` works on `:cpu` and `:gpu`. +`backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD. +`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. # Internal API @@ -355,8 +360,8 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) :wd => 0.0f0, :batchsize => 2048, :seed => 123, - :backend => :enzyme, - :device => :cpu, + :backend => :zygote, + :device => :gpu, :gpuID => 0, :embedding_config => nothing, ) @@ -378,8 +383,11 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) backend = Symbol(args[:backend]) device = Symbol(args[:device]) - if device == :reactant && backend != :enzyme - error("`device=:reactant` requires `backend=:enzyme`. Got `backend=:$backend`.") + if device == :reactant + error("Use `backend=:reactant` with `device=:cpu` or `device=:gpu` instead of `device=:reactant`.") + end + if backend == :enzyme && device == :gpu + @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU." end # Build EmbeddingConfig from Dict if needed diff --git a/test/core.jl b/test/core.jl index 7605d67..b89d29c 100644 --- a/test/core.jl +++ b/test/core.jl @@ -210,7 +210,7 @@ end end -@testset "Backend/device - reactant requires enzyme" begin +@testset "Backend/device - reactant is a backend" begin @test_throws ErrorException NeuroTabRegressor(; backend=:zygote, device=:reactant) @test_throws ErrorException NeuroTabClassifier(; backend=:zygote, device=:reactant) end @@ -218,7 +218,7 @@ end @testset "Backend/device - Regression ($backend, $device)" for (backend, device) in [ (:enzyme, :cpu), (:zygote, :cpu), - (:enzyme, :reactant), + (:reactant, :cpu), ] Random.seed!(123) From 90a110a37129e0a660e61d46a391e553a6ec61a0 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 19:48:13 -0400 Subject: [PATCH 100/120] up --- Project.toml | 1 + benchmarks/Project.toml | 18 +++++++++++ benchmarks/benchmark_mse.jl | 57 ++++++++++++++++++++--------------- benchmarks/titanic-logloss.jl | 9 ++++-- 4 files changed, 59 insertions(+), 26 deletions(-) create mode 100644 benchmarks/Project.toml diff --git a/Project.toml b/Project.toml index 3ee8dbc..46ad5e0 100644 --- a/Project.toml +++ b/Project.toml @@ -36,6 +36,7 @@ Reactant = "0.2.247" Statistics = "1" StatsBase = "0.34" Tables = "1.9" +Zygote = "0.7" julia = "1.10" [weakdeps] diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml new file mode 100644 index 0000000..3b6c210 --- /dev/null +++ b/benchmarks/Project.toml @@ -0,0 +1,18 @@ +[deps] +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" +DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" +NeuroTabModels = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" + +[compat] +CUDA = "5" diff --git a/benchmarks/benchmark_mse.jl b/benchmarks/benchmark_mse.jl index 28b3772..2ad2646 100644 --- a/benchmarks/benchmark_mse.jl +++ b/benchmarks/benchmark_mse.jl @@ -3,6 +3,11 @@ using DataFrames using BenchmarkTools using Random: seed! +using CUDA, cuDNN +using Reactant +using Zygote +using Enzyme + Threads.nthreads() seed!(123) @@ -16,28 +21,28 @@ feature_names = names(dtrain) dtrain.y = Y target_name = "y" -# arch = NeuroTabModels.NeuroTreeConfig(; -# tree_type=:binary, -# actA=:identity, -# init_scale=1.0, -# depth=4, -# ntrees=32, -# stack_size=1, -# hidden_size=1, -# scaler=false, -# ) -arch = NeuroTabModels.TabMConfig(; - arch_type=:tabm, - k=16, - d_block=64, - n_blocks=3, - dropout=0.1, - bins=nothing, - use_embeddings=false, - embedding_type=:periodic, - d_embedding=16, - scaling_init=:random_signs, +arch = NeuroTabModels.NeuroTreeConfig(; + tree_type=:binary, + actA=:identity, + init_scale=1.0, + depth=4, + ntrees=32, + stack_size=1, + hidden_size=1, + scaler=false, ) +# arch = NeuroTabModels.TabMConfig(; +# arch_type=:tabm, +# k=16, +# d_block=64, +# n_blocks=3, +# dropout=0.1, +# bins=nothing, +# use_embeddings=false, +# embedding_type=:periodic, +# d_embedding=16, +# scaling_init=:random_signs, +# ) # arch = NeuroTabModels.MLPConfig(; # act=:relu, # stack_size=1, @@ -50,12 +55,16 @@ learner = NeuroTabRegressor( nrounds=10, lr=1e-2, batchsize=2048, - device=:gpu + device=:gpu, + backend=:enzyme ) # Reactant GPU: 5.970480 seconds (2.33 M allocations: 5.242 GiB, 3.80% gc time, 0.00% compilation time) -# Zygote GPU: 9.855853 seconds (27.92 M allocations: 6.005 GiB, 3.58% gc time) -# 13.557744 seconds (26.40 M allocations: 5.989 GiB, 9.60% gc time) +# Reactant GPU with eval: 10.154589 seconds (2.33 M allocations: 10.563 GiB, 17.66% gc time, 0.00% compilation time: 100% of which was recompilation) +# Zygote GPU: 8.798624 seconds (15.09 M allocations: 5.728 GiB, 13.60% gc time) +# Zygote GPU with eval: 13.715713 seconds (20.61 M allocations: 11.236 GiB, 22.54% gc time) +# Zygote CPU: 338.912009 seconds (62.93 M allocations: 373.382 GiB, 10.93% gc time, 6.38% compilation time: <1% of which was recompilation) +# Enzyme CPU: 657.713208 seconds (1.98 M allocations: 270.987 GiB, 6.37% gc time) @time m = NeuroTabModels.fit( learner, dtrain; diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 69366cc..c6d8cf1 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -2,12 +2,16 @@ using MLDatasets using DataFrames using Statistics: mean using StatsBase: median -using CategoricalArrays using Random using CategoricalArrays using OrderedCollections using NeuroTabModels +using CUDA, cuDNN +using Reactant +using Zygote +using Enzyme + Random.seed!(123) df = MLDatasets.Titanic().dataframe @@ -78,7 +82,8 @@ learner = NeuroTabRegressor( nrounds=200, early_stopping_rounds=2, lr=1e-2, - device=:gpu + device=:reactant, + backend=:enzyme ) @time m = NeuroTabModels.fit( From e38782b1072c1a99778774307defa85b47c29908 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 20:36:06 -0400 Subject: [PATCH 101/120] up --- benchmarks/Project.toml | 3 ++- benchmarks/YEAR-regression.jl | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 3b6c210..027cbbc 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -15,4 +15,5 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] -CUDA = "5" +CUDA = "6" +cuDNN = "6" diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index ba2a2be..9c529e0 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -7,6 +7,11 @@ using NeuroTabModels using AWS: AWSCredentials, AWSConfig, @service @service S3 +using Enzyme +using Reactant +using Zygote +using CUDA, cuDNN + aws_creds = AWSCredentials(ENV["AWS_ACCESS_KEY_ID_JDB"], ENV["AWS_SECRET_ACCESS_KEY_JDB"]) aws_config = AWSConfig(; creds=aws_creds, region="ca-central-1") @@ -84,6 +89,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # ) device = :gpu +backend = :zygote loss = :mse # :mse :gaussian_mle :tweedie # embedding_config = Dict( From 81b0279a043e89616581189cc7de8b59d36c3c4e Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 21:47:52 -0400 Subject: [PATCH 102/120] up --- adhoc/Project.toml | 5 +++++ benchmarks/Project.toml | 5 +++-- benchmarks/YEAR-regression.jl | 14 ++++++++------ cuda-debug/Project.toml | 10 ++++++++++ 4 files changed, 26 insertions(+), 8 deletions(-) create mode 100644 adhoc/Project.toml create mode 100644 cuda-debug/Project.toml diff --git a/adhoc/Project.toml b/adhoc/Project.toml new file mode 100644 index 0000000..a0cc115 --- /dev/null +++ b/adhoc/Project.toml @@ -0,0 +1,5 @@ +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" +cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 027cbbc..6ab1d45 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -11,9 +11,10 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] -CUDA = "6" -cuDNN = "6" +CUDA = "5.11" +cuDNN = "1" \ No newline at end of file diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 9c529e0..3d1e61c 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -3,14 +3,15 @@ using CSV using DataFrames using Statistics: mean, std using StatsBase: tiedrank -using NeuroTabModels -using AWS: AWSCredentials, AWSConfig, @service -@service S3 +using CUDA, cuDNN using Enzyme using Reactant using Zygote -using CUDA, cuDNN + +using NeuroTabModels +using AWS: AWSCredentials, AWSConfig, @service +@service S3 aws_creds = AWSCredentials(ENV["AWS_ACCESS_KEY_ID_JDB"], ENV["AWS_SECRET_ACCESS_KEY_JDB"]) aws_config = AWSConfig(; creds=aws_creds, region="ca-central-1") @@ -89,7 +90,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # ) device = :gpu -backend = :zygote +backend = :reactant loss = :mse # :mse :gaussian_mle :tweedie # embedding_config = Dict( @@ -109,7 +110,8 @@ learner = NeuroTabRegressor( early_stopping_rounds=2, lr=1e-3, batchsize=512, - device + device, + backend ) @time m = NeuroTabModels.fit( diff --git a/cuda-debug/Project.toml b/cuda-debug/Project.toml new file mode 100644 index 0000000..0f718a0 --- /dev/null +++ b/cuda-debug/Project.toml @@ -0,0 +1,10 @@ +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +NeuroTabModels = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" +cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" + +[compat] +CUDA = "5" From 15128741b341072fd04b5638f40652cec7433ab7 Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Sun, 3 May 2026 07:19:14 +0530 Subject: [PATCH 103/120] Multi-AD backend support via package extensions (#44) * add support for multi-AD backend through extension Signed-off-by: AdityaPandeyCN * revert default to cpu Signed-off-by: AdityaPandeyCN * REMOVE ZYGOTE COMPAT Signed-off-by: AdityaPandeyCN * Move Reactant selection from to , keep Enzyme as the AD backend underneath, and default learners to Zygote/GPU for broader compatibility. Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN --- Project.toml | 20 +++++-- ext/NeuroTabModelsEnzymeExt.jl | 11 ++++ ext/NeuroTabModelsReactantExt.jl | 50 +++++++++++++++++ ext/NeuroTabModelsZygoteExt.jl | 11 ++++ src/Fit/callback.jl | 19 ++++--- src/Fit/fit.jl | 38 +++++++------ src/infer.jl | 89 ++++++++++++++++------------- src/learners.jl | 40 +++++++++++-- test/core.jl | 96 ++++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 10 files changed, 301 insertions(+), 76 deletions(-) create mode 100644 ext/NeuroTabModelsEnzymeExt.jl create mode 100644 ext/NeuroTabModelsReactantExt.jl create mode 100644 ext/NeuroTabModelsZygoteExt.jl diff --git a/Project.toml b/Project.toml index 46f0759..61e6f2e 100644 --- a/Project.toml +++ b/Project.toml @@ -6,7 +6,6 @@ authors = ["jeremie.db "] [deps] CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623" LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11" @@ -16,7 +15,6 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f" Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" @@ -34,16 +32,30 @@ NNlib = "0.9" OneHotArrays = "0.2" Optimisers = "0.4" Random = "1" -Reactant = "0.2.247" +Reactant = "0.2.254" Statistics = "1" StatsBase = "0.34" Tables = "1.9" +Zygote = "0.7" julia = "1.10" +[weakdeps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[extensions] +NeuroTabModelsEnzymeExt = "Enzyme" +NeuroTabModelsReactantExt = "Reactant" +NeuroTabModelsZygoteExt = "Zygote" + [extras] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" MLJBase = "a7f614a8-145f-11e9-1d2a-a57a1082229d" MLJTestInterface = "72560011-54dd-4dc2-94f3-c5de45b75ecd" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Test", "MLJTestInterface", "MLJBase"] +test = ["Test", "MLJTestInterface", "MLJBase", "Enzyme", "Reactant", "Zygote"] diff --git a/ext/NeuroTabModelsEnzymeExt.jl b/ext/NeuroTabModelsEnzymeExt.jl new file mode 100644 index 0000000..9ba5b1e --- /dev/null +++ b/ext/NeuroTabModelsEnzymeExt.jl @@ -0,0 +1,11 @@ +module NeuroTabModelsEnzymeExt + +using Enzyme +using Lux: AutoEnzyme +using NeuroTabModels + +import NeuroTabModels.Fit: get_ad_backend + +get_ad_backend(::Val{:enzyme}) = AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse)) + +end diff --git a/ext/NeuroTabModelsReactantExt.jl b/ext/NeuroTabModelsReactantExt.jl new file mode 100644 index 0000000..791cc74 --- /dev/null +++ b/ext/NeuroTabModelsReactantExt.jl @@ -0,0 +1,50 @@ +module NeuroTabModelsReactantExt + +using Lux: Training, reactant_device +using NeuroTabModels +using Reactant +using Reactant: @compile + +import NeuroTabModels.Infer: _get_device, _infer_loop, _infer_grp_loop +import NeuroTabModels.Fit: _single_train_step! +import NeuroTabModels.Fit.CallBacks: _compile_eval_step + +using NeuroTabModels.Infer: _forward_reduce + +function _get_device(::Val{:reactant}, ::Val{D}; gpuID::Integer=0) where {D} + Reactant.set_default_backend(String(D)) + return reactant_device() +end + +function _infer_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st) + compiled = @compile _forward_reduce(chain, dev(x0), ps, st) + + preds = Vector{AbstractArray}() + for x in data + if size(x) == size(x0) + pred = compiled(chain, dev(x), ps, st) + else + pred = Reactant.@jit _forward_reduce(chain, dev(x), ps, st) + end + push!(preds, cdev(pred)) + end + return preds +end + +function _infer_grp_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st) + compiled = @compile _forward_reduce(chain, dev(x0), ps, st) + + preds = Vector{AbstractArray}() + for (x, mask) in data + pred = compiled(chain, dev(x), ps, st) + push!(preds, cdev(pred)[:, mask]) + end + return preds +end + +_single_train_step!(::Val{:reactant}, ad_backend, lux_loss, d, ts) = + Training.single_train_step!(ad_backend, lux_loss, d, ts; return_gradients=Val(false)) + +_compile_eval_step(::Val{:reactant}, step, args...) = @compile step(args...) + +end diff --git a/ext/NeuroTabModelsZygoteExt.jl b/ext/NeuroTabModelsZygoteExt.jl new file mode 100644 index 0000000..b5556f3 --- /dev/null +++ b/ext/NeuroTabModelsZygoteExt.jl @@ -0,0 +1,11 @@ +module NeuroTabModelsZygoteExt + +using Lux: AutoZygote +using NeuroTabModels +using Zygote + +import NeuroTabModels.Fit: get_ad_backend + +get_ad_backend(::Val{:zygote}) = AutoZygote() + +end diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index e16de91..04b7f7b 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -4,12 +4,11 @@ using DataFrames using Statistics: mean, median using ..Learners: LearnerTypes -using ...Infer: reduce_pred +using ...Infer: reduce_pred, _get_device using ..Data: get_df_loader_train using ..Metrics -using Lux: Training, reactant_device, testmode -using Reactant: @compile, set_default_backend +using Lux: Training, testmode export CallBack, init_logger, update_logger!, agg_logger @@ -34,7 +33,7 @@ function CallBack( offset_name=nothing, group_key=nothing ) - dev = reactant_device() + dev = _get_device(config.backend, config.device; gpuID=config.gpuID) ts = cache[:train_state] scalers = cache[:scalers] batchsize = config.batchsize @@ -45,33 +44,35 @@ function CallBack( ps, st = ts.parameters, testmode(ts.states) d0 = first(deval) - eval_compiled = _compile_eval_step(ts.model, feval, d0, ps, st) + eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.backend == :reactant) return CallBack(deval, eval_compiled) end -function _compile_eval_step(chain, feval, d0, ps, st) +function _build_eval_step(chain, feval, d0, ps, st; reactant::Bool) if length(d0) == 2 function _step2(x, y, ps, st) m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y; agg=sum), eltype(y)(last(size(y))) end - return @compile _step2(d0[1], d0[2], ps, st) + return reactant ? _compile_eval_step(Val(:reactant), _step2, d0[1], d0[2], ps, st) : _step2 elseif length(d0) == 3 function _step3(x, y, w, ps, st) m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y, w; agg=sum), sum(w) end - return @compile _step3(d0[1], d0[2], d0[3], ps, st) + return reactant ? _compile_eval_step(Val(:reactant), _step3, d0[1], d0[2], d0[3], ps, st) : _step3 else function _step4(x, y, w, offset, ps, st) m = x -> reduce_pred(first(chain(x, ps, st))) return feval(m, x, y, w, offset; agg=sum), sum(w) end - return @compile _step4(d0[1], d0[2], d0[3], d0[4], ps, st) + return reactant ? _compile_eval_step(Val(:reactant), _step4, d0[1], d0[2], d0[3], d0[4], ps, st) : _step4 end end +_compile_eval_step(::Val{D}, step, args...) where {D} = step + function init_logger(config::LearnerTypes) logger = Dict( :name => String(config.metric), diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index b925c3f..cf9e41b 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -7,33 +7,28 @@ using ..Learners using ..Models using ..Losses using ..Metrics -using ..Infer: reduce_pred +using ..Infer: reduce_pred, _get_device import Random: Xoshiro import Statistics: mean, std import MLJModelInterface: fit import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam using Lux -using Reactant -using Lux: cpu_device, gpu_device, reactant_device +using Lux: cpu_device using DataFrames using CategoricalArrays +get_ad_backend(backend::Symbol) = get_ad_backend(Val(backend)) +get_ad_backend(::Val{:reactant}) = get_ad_backend(Val(:enzyme)) +get_ad_backend(::Val{b}) where {b} = error( + "Unsupported or unloaded `backend=:$b`. Supported: [:enzyme, :zygote, :reactant]. " * + "`:enzyme` and `:reactant` require `using Enzyme`, `:zygote` requires `using Zygote`." +) + include("callback.jl") using .CallBacks -""" - _get_device(device::Symbol) - -!!! warning - Returns the default reactant device. - Future behavior should support :cpu/gpu device in addition to the assumed :reactant device. -""" -function _get_device(device::Symbol) - return reactant_device() -end - function init( config::LearnerTypes, df::AbstractDataFrame; @@ -49,7 +44,7 @@ function init( offset_name = isnothing(offset_name) ? nothing : Symbol(offset_name) group_key = isnothing(group_key) ? nothing : Symbol(group_key) - dev = _get_device(config.device) + dev = _get_device(config.backend, config.device; gpuID=config.gpuID) batchsize = config.batchsize nfeats = length(feature_names) L = get_loss_type(config.loss) @@ -101,7 +96,10 @@ function init( :group_key => group_key, :target_levels => target_levels, :target_isordered => target_isordered, - :scalers => scalers + :scalers => scalers, + :backend => config.backend, + :device => config.device, + :gpuID => config.gpuID ) m = NeuroTabModel(L, chain, info) @@ -110,7 +108,7 @@ function init( opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) - return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts, :scalers => scalers) + return m, Dict(:data => data, :lux_loss => lux_loss, :train_state => ts, :scalers => scalers, :ad_backend => get_ad_backend(config.backend)) end """ @@ -201,10 +199,14 @@ function _sync_params_to_model!(m, cache) m.info[:st] = cdev(Lux.testmode(ts.states)) end +_single_train_step!(device::Symbol, ad_backend, lux_loss, d, ts) = _single_train_step!(Val(device), ad_backend, lux_loss, d, ts) +_single_train_step!(::Val{D}, ad_backend, lux_loss, d, ts) where {D} = Training.single_train_step!(ad_backend, lux_loss, d, ts) + function fit_iter!(m, cache) ts, lux_loss = cache[:train_state], cache[:lux_loss] + ad_backend = cache[:ad_backend] for d in cache[:data] - _, loss, _, ts = Training.single_train_step!(AutoEnzyme(), lux_loss, d, ts) + _, loss, _, ts = _single_train_step!(m.info[:backend], ad_backend, lux_loss, d, ts) end cache[:train_state] = ts m.info[:nrounds] += 1 diff --git a/src/infer.jl b/src/infer.jl index cc1939d..8bb0089 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -5,9 +5,7 @@ using ..Losses using ..Models using Lux -using Lux: cpu_device, gpu_device, reactant_device -using Reactant -using Reactant: @compile +using Lux: cpu_device, gpu_device using NNlib: sigmoid, softmax using Statistics: mean using DataFrames: AbstractDataFrame, GroupedDataFrame, groupby @@ -18,17 +16,20 @@ export infer, reduce_pred reduce_pred(pred::AbstractMatrix) = pred reduce_pred(pred::AbstractArray{T,3}) where {T} = dropdims(mean(pred; dims=2); dims=2) - """ - _get_device(device::Symbol) + _get_device(device::Symbol; gpuID::Integer=0) -!!! warning - Returns the default reactant device. - Future behavior should support :cpu/gpu device in addition to the assumed :reactant device. +Resolve a Lux device from a symbol: `:cpu` or `:gpu`. +`backend=:reactant` requires `using Reactant` to activate `NeuroTabModelsReactantExt`. """ -function _get_device(device::Symbol) - return reactant_device() -end +_get_device(device::Symbol; gpuID::Integer=0) = _get_device(Val(device); gpuID) +_get_device(backend::Symbol, device::Symbol; gpuID::Integer=0) = _get_device(Val(backend), Val(device); gpuID) +_get_device(::Val, ::Val{D}; gpuID::Integer=0) where {D} = _get_device(Val(D); gpuID) +_get_device(::Val{:cpu}; gpuID::Integer=0) = cpu_device() +_get_device(::Val{:gpu}; gpuID::Integer=0) = gpu_device(gpuID == 0 ? nothing : gpuID) +_get_device(::Val{D}; gpuID::Integer=0) where {D} = error( + "Unsupported `device=:$D`. Supported devices are [:cpu, :gpu]." +) function _forward_reduce(chain, x, ps, st) pred, _ = chain(x, ps, st) @@ -60,25 +61,38 @@ function _scaler(::Type{<:GaussianMLE}, p, scalers::NamedTuple) return p end -function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} - dev = _get_device(Symbol(device)) +# Default inference loops for :cpu and :gpu. +# NeuroTabModelsReactantExt adds ::Val{:reactant} methods that run a Reactant-compiled forward pass. +function _infer_loop(::Val, chain, data, x0, dev, cdev, ps, st) + preds = Vector{AbstractArray}() + for x in data + pred = _forward_reduce(chain, dev(x), ps, st) + push!(preds, cdev(pred)) + end + return preds +end + +# Grouped variant: each batch is `(x, mask)`; `mask` selects real rows from padded groups. +function _infer_grp_loop(::Val, chain, data, x0, dev, cdev, ps, st) + preds = Vector{AbstractArray}() + for (x, mask) in data + pred = _forward_reduce(chain, dev(x), ps, st) + push!(preds, cdev(pred)[:, mask]) + end + return preds +end + +function infer(m::NeuroTabModel{L}, data; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) where {L} + device = Symbol(device) + backend = Symbol(backend) + dev = _get_device(backend, device; gpuID=get(m.info, :gpuID, 0)) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) scalers = m.info[:scalers] x0 = first(data) - compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) - - preds = Vector{AbstractArray}() - for x in data - if size(x) == size(x0) - pred = compiled(m.chain, dev(x), ps, st) - else - pred = Reactant.@jit _forward_reduce(m.chain, dev(x), ps, st) - end - push!(preds, cdev(pred)) - end + preds = _infer_loop(Val(backend), m.chain, data, x0, dev, cdev, ps, st) p_raw = _assemble(L, preds) proj || return p_raw @@ -86,8 +100,10 @@ function infer(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L return _scaler(L, p, scalers) end -function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) where {L} - dev = _get_device(Symbol(device)) +function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) where {L} + device = Symbol(device) + backend = Symbol(backend) + dev = _get_device(backend, device; gpuID=get(m.info, :gpuID, 0)) cdev = cpu_device() ps = dev(m.info[:ps]) st = dev(m.info[:st]) @@ -98,13 +114,7 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher # data = data |> dev # (x0, mask0) = first(data) # @info typeof("mask0-dev") mask0 - compiled = @compile _forward_reduce(m.chain, dev(x0), ps, st) - - preds = Vector{AbstractArray}() - for (x, mask) in data - pred = compiled(m.chain, dev(x), ps, st) - push!(preds, cdev(pred)[:, mask]) - end + preds = _infer_grp_loop(Val(backend), m.chain, data, x0, dev, cdev, ps, st) p_raw = _assemble(L, preds) proj || return p_raw @@ -112,25 +122,26 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, proj::Bool=true) wher return _scaler(L, p, scalers) end -function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, proj::Bool=true) +function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) group_key = m.info[:group_key] if isnothing(group_key) dinfer = get_df_loader_infer(df; feature_names=m.info[:feature_names], batchsize=2048) - p = infer(m, dinfer; device, proj) + p = infer(m, dinfer; device, backend, proj) else dfg = groupby(df, group_key; sort=true) dinfer = get_df_loader_infer(dfg; feature_names=m.info[:feature_names], batchsize=2048) - p = infer_grp(m, dinfer; device, proj) + p = infer_grp(m, dinfer; device, backend, proj) end return p end -function (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, proj::Bool=true) - return infer(m, df; device, proj) +function (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) + return infer(m, df; device, backend, proj) end # function (m::NeuroTabModel)(x::AbstractMatrix; device=:cpu) # return infer(m, [(x,)]; device) # end -end \ No newline at end of file +end + diff --git a/src/learners.jl b/src/learners.jl index 37c254f..3e4d8d9 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -20,6 +20,7 @@ mutable struct NeuroTabRegressor <: MMI.Deterministic batchsize::Int seed::Int scale_target::Bool + backend::Symbol device::Symbol gpuID::Int end @@ -43,11 +44,15 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl]( - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. -- `device=:cpu`: Device on which to perform the computation, either `:cpu` or `:gpu` -- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `backend=:zygote`: Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`. +- `device=:gpu`: Execution device. One of `:cpu` or `:gpu`. +- `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`. +`backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD. +`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. + # Internal API Do `config = NeuroTabRegressor()` to construct an instance with default hyper-parameters. @@ -145,7 +150,8 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) :wd => 0.0f0, :batchsize => 2048, :seed => 123, - :device => :cpu, + :backend => :zygote, + :device => :gpu, :gpuID => 0, :embedding_config => nothing, :scale_target => true @@ -179,7 +185,14 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) error("Invalid metric. Must be one of: $_metric_list") end + backend = Symbol(args[:backend]) device = Symbol(args[:device]) + if device == :reactant + error("Use `backend=:reactant` with `device=:cpu` or `device=:gpu` instead of `device=:reactant`.") + end + if backend == :enzyme && device == :gpu + @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU." + end # Build EmbeddingConfig from Dict if needed embed = args[:embedding_config] @@ -199,6 +212,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) args[:batchsize], args[:seed], args[:scale_target], + backend, device, args[:gpuID] ) @@ -224,6 +238,7 @@ mutable struct NeuroTabClassifier <: MMI.Probabilistic wd::Float32 batchsize::Int seed::Int + backend::Symbol device::Symbol gpuID::Int end @@ -241,10 +256,14 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl] - `wd=0.f0`: Weight decay applied to the gradients by the optimizer. - `batchsize=2048`: Batch size. - `seed=123`: An integer used as a seed to the random number generator. -- `device=:cpu`: Device on which to perform the computation, either `:cpu` or `:gpu` -- `gpuID=0`: GPU device to use, only relveant if `device = :gpu` +- `backend=:zygote`: Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`. +- `device=:gpu`: Execution device. One of `:cpu` or `:gpu`. +- `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. - `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. +`backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD. +`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. + # Internal API Do `config = NeuroTabClassifier()` to construct an instance with default hyper-parameters. @@ -341,7 +360,8 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) :wd => 0.0f0, :batchsize => 2048, :seed => 123, - :device => :cpu, + :backend => :zygote, + :device => :gpu, :gpuID => 0, :embedding_config => nothing, ) @@ -361,7 +381,14 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) args[arg] = kwargs[arg] end + backend = Symbol(args[:backend]) device = Symbol(args[:device]) + if device == :reactant + error("Use `backend=:reactant` with `device=:cpu` or `device=:gpu` instead of `device=:reactant`.") + end + if backend == :enzyme && device == :gpu + @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU." + end # Build EmbeddingConfig from Dict if needed embed = args[:embedding_config] @@ -380,6 +407,7 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) Float32(args[:wd]), args[:batchsize], args[:seed], + backend, device, args[:gpuID] ) diff --git a/test/core.jl b/test/core.jl index eb4c74a..b89d29c 100644 --- a/test/core.jl +++ b/test/core.jl @@ -208,4 +208,100 @@ end @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 @test mean(peval .== levelcode.(deval.class)) > 0.95 +end + +@testset "Backend/device - reactant is a backend" begin + @test_throws ErrorException NeuroTabRegressor(; backend=:zygote, device=:reactant) + @test_throws ErrorException NeuroTabClassifier(; backend=:zygote, device=:reactant) +end + +@testset "Backend/device - Regression ($backend, $device)" for (backend, device) in [ + (:enzyme, :cpu), + (:zygote, :cpu), + (:reactant, :cpu), +] + + Random.seed!(123) + X = randn(Float32, 1000, 10) + y = X[:, 1] .+ 0.5f0 .* X[:, 2] .+ 0.1f0 .* randn(Float32, 1000) + df = DataFrame(X, :auto) + df[!, :y] = y + target_name = "y" + feature_names = setdiff(names(df), [target_name]) + + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + learner = NeuroTabRegressor(; + arch_name="NeuroTreeConfig", + arch_config=Dict(:depth => 3), + loss=:mse, + nrounds=20, + early_stopping_rounds=2, + lr=1e-1, + backend, + device, + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + target_name, + feature_names, + deval, + ) + + p = m(deval) + @test size(p, 1) == nrow(deval) + @test !any(isnan, p) + mse_model = mean((p .- deval.y) .^ 2) + mse_baseline = mean((mean(dtrain.y) .- deval.y) .^ 2) + @test mse_model < mse_baseline +end + +@testset "Backend/device - Classification ($backend, $device)" for (backend, device) in [ + (:enzyme, :cpu), + (:zygote, :cpu), +] + + Random.seed!(123) + X, y = @load_crabs + df = DataFrame(X) + df[!, :class] = y + target_name = "class" + feature_names = setdiff(names(df), [target_name]) + + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + learner = NeuroTabClassifier(; + arch_name="NeuroTreeConfig", + arch_config=Dict(:depth => 4), + embedding_config=Dict(:embedding_type => :batchnorm), + nrounds=200, + early_stopping_rounds=5, + lr=3e-2, + backend, + device, + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + deval, + target_name, + feature_names, + ) + + ptrain = [argmax(x) for x in eachrow(m(dtrain))] + peval = [argmax(x) for x in eachrow(m(deval))] + @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 + @test mean(peval .== levelcode.(deval.class)) > 0.95 + end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 8959b2b..a60810b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,6 +8,9 @@ using StatsBase: sample using Random using MLJBase using MLJTestInterface +using Enzyme +using Reactant +using Zygote include("core.jl") include("embedding.jl") From 8b4cb2ebcaccb7aeb1fc837593246efb219faabf Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 22:14:13 -0400 Subject: [PATCH 104/120] up --- src/metrics.jl | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/src/metrics.jl b/src/metrics.jl index ea9e3c4..19222bb 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -140,6 +140,45 @@ function gaussian_mle(m, x, y, w, offset; agg=mean) return metric end + +""" + correlation(m, x, y; agg=mean) + correlation(m, x, y, w; agg=mean) + correlation(m, x, y, w, offset; agg=mean) +""" +function correlation(m, x, y; agg=mean) + p = exp.(vec(m(x))) + y = vec(y) + p_mean = mean(p) + p_var = mean(p .^ 2) - p_mean^2 + y_mean = mean(y) + y_var = mean(y .^ 2) - y_mean^2 + py_mean = mean(p .* y) + return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * length(y) +end +function correlation(m, x, y, w; agg=mean) + p = exp.(vec(m(x))) + y = vec(y) + w = vec(w) + p_mean = w' * p + p_var = w' * (p .^ 2) - p_mean^2 + y_mean = w' * y + y_var = w' * (y .^ 2) - y_mean^2 + py_mean = w' * (p .* y) + return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) +end +function correlation(m, x, y, w; agg=mean) + p = exp.(vec(m(x)) .+ vec(offset)) + y = vec(y) + w = vec(w) + p_mean = w' * p + p_var = w' * (p .^ 2) - p_mean^2 + y_mean = w' * y + y_var = w' * (y .^ 2) - y_mean^2 + py_mean = w' * (p .* y) + return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) +end + function get_metric(ts, data, eval_compiled) metric = 0.0f0 ws = 0.0f0 @@ -165,6 +204,7 @@ const metric_dict = Dict( :mlogloss => mlogloss, :gaussian_mle => gaussian_mle, :tweedie => tweedie, + :correlation => correlation, ) is_maximise(::typeof(mse)) = false @@ -173,5 +213,6 @@ is_maximise(::typeof(logloss)) = false is_maximise(::typeof(mlogloss)) = false is_maximise(::typeof(gaussian_mle)) = true is_maximise(::typeof(tweedie)) = false +is_maximise(::typeof(correlation)) = true end \ No newline at end of file From ae427c4bc674d82fab5dff81f23616cc9097e494 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 22:27:56 -0400 Subject: [PATCH 105/120] correlation metric --- adhoc/Project.toml | 5 ----- benchmarks/YEAR-regression.jl | 2 ++ cuda-debug/Project.toml | 10 ---------- src/learners.jl | 2 +- src/metrics.jl | 9 ++++++--- 5 files changed, 9 insertions(+), 19 deletions(-) delete mode 100644 adhoc/Project.toml delete mode 100644 cuda-debug/Project.toml diff --git a/adhoc/Project.toml b/adhoc/Project.toml deleted file mode 100644 index a0cc115..0000000 --- a/adhoc/Project.toml +++ /dev/null @@ -1,5 +0,0 @@ -[deps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -Lux = "b2108857-7c20-44ae-9111-449ecde12c47" -WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d" -cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 3d1e61c..af30ac4 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -92,6 +92,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; device = :gpu backend = :reactant loss = :mse # :mse :gaussian_mle :tweedie +metric = :correlation # :mse :gaussian_mle :tweedie # embedding_config = Dict( # :embedding_type => :piecewise, @@ -106,6 +107,7 @@ learner = NeuroTabRegressor( arch; embedding_config, loss, + metric, nrounds=200, early_stopping_rounds=2, lr=1e-3, diff --git a/cuda-debug/Project.toml b/cuda-debug/Project.toml deleted file mode 100644 index 0f718a0..0000000 --- a/cuda-debug/Project.toml +++ /dev/null @@ -1,10 +0,0 @@ -[deps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" -NeuroTabModels = "f03403ce-56d7-46f9-9b5e-ff6add8ca7b3" -Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" -Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" -cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" - -[compat] -CUDA = "5" diff --git a/src/learners.jl b/src/learners.jl index 3e4d8d9..811e00a 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -175,7 +175,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) loss = Symbol(args[:loss]) loss ∉ [:mse, :mae, :logloss, :tweedie, :gaussian_mle] && error("The provided kwarg `loss`: $loss is not supported.") - _metric_list = [:mse, :mae, :logloss, :tweedie, :gaussian_mle] + _metric_list = [:mse, :mae, :logloss, :tweedie, :gaussian_mle, :correlation] if isnothing(args[:metric]) metric = loss else diff --git a/src/metrics.jl b/src/metrics.jl index 19222bb..883f738 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -147,7 +147,8 @@ end correlation(m, x, y, w, offset; agg=mean) """ function correlation(m, x, y; agg=mean) - p = exp.(vec(m(x))) + _p = exp.(vec(m(x))) + p = view(_p, 1, :) y = vec(y) p_mean = mean(p) p_var = mean(p .^ 2) - p_mean^2 @@ -157,7 +158,8 @@ function correlation(m, x, y; agg=mean) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * length(y) end function correlation(m, x, y, w; agg=mean) - p = exp.(vec(m(x))) + _p = exp.(vec(m(x))) + p = view(_p, 1, :) y = vec(y) w = vec(w) p_mean = w' * p @@ -168,7 +170,8 @@ function correlation(m, x, y, w; agg=mean) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) end function correlation(m, x, y, w; agg=mean) - p = exp.(vec(m(x)) .+ vec(offset)) + _p = exp.(vec(m(x)) .+ vec(offset)) + p = view(_p, 1, :) y = vec(y) w = vec(w) p_mean = w' * p From 9e76eead983838aeae4e5ee0171ce04f9e2312d2 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 22:37:51 -0400 Subject: [PATCH 106/120] correlation metric --- src/metrics.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/metrics.jl b/src/metrics.jl index 883f738..1613bc6 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -147,7 +147,7 @@ end correlation(m, x, y, w, offset; agg=mean) """ function correlation(m, x, y; agg=mean) - _p = exp.(vec(m(x))) + _p = vec(m(x)) p = view(_p, 1, :) y = vec(y) p_mean = mean(p) @@ -158,7 +158,7 @@ function correlation(m, x, y; agg=mean) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * length(y) end function correlation(m, x, y, w; agg=mean) - _p = exp.(vec(m(x))) + _p = vec(m(x)) p = view(_p, 1, :) y = vec(y) w = vec(w) @@ -170,7 +170,7 @@ function correlation(m, x, y, w; agg=mean) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) end function correlation(m, x, y, w; agg=mean) - _p = exp.(vec(m(x)) .+ vec(offset)) + _p = vec(m(x)) .+ vec(offset) p = view(_p, 1, :) y = vec(y) w = vec(w) From c5182939b47be48d20d8bc59aa2fa230001a6358 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 22:55:35 -0400 Subject: [PATCH 107/120] correlation metric --- src/metrics.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/metrics.jl b/src/metrics.jl index 1613bc6..75f2dfc 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -147,8 +147,7 @@ end correlation(m, x, y, w, offset; agg=mean) """ function correlation(m, x, y; agg=mean) - _p = vec(m(x)) - p = view(_p, 1, :) + p = view(m(x), 1, :) y = vec(y) p_mean = mean(p) p_var = mean(p .^ 2) - p_mean^2 @@ -158,8 +157,7 @@ function correlation(m, x, y; agg=mean) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * length(y) end function correlation(m, x, y, w; agg=mean) - _p = vec(m(x)) - p = view(_p, 1, :) + p = view(m(x), 1, :) y = vec(y) w = vec(w) p_mean = w' * p @@ -169,9 +167,8 @@ function correlation(m, x, y, w; agg=mean) py_mean = w' * (p .* y) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) end -function correlation(m, x, y, w; agg=mean) - _p = vec(m(x)) .+ vec(offset) - p = view(_p, 1, :) +function correlation(m, x, y, w, offset; agg=mean) + p = view(m(x), 1, :) .+ view(offset, 1, :) y = vec(y) w = vec(w) p_mean = w' * p From 1fc5fbacaf1f0b338df35e1d5a65e4e65c098f6e Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 2 May 2026 23:21:09 -0400 Subject: [PATCH 108/120] up --- benchmarks/YEAR-regression.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index af30ac4..569a876 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -92,7 +92,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; device = :gpu backend = :reactant loss = :mse # :mse :gaussian_mle :tweedie -metric = :correlation # :mse :gaussian_mle :tweedie +# metric = :correlation # :mse :gaussian_mle :tweedie # embedding_config = Dict( # :embedding_type => :piecewise, @@ -107,7 +107,7 @@ learner = NeuroTabRegressor( arch; embedding_config, loss, - metric, + # metric, nrounds=200, early_stopping_rounds=2, lr=1e-3, From d41074772494c0318933c5d2b14e592282916816 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Sun, 3 May 2026 16:42:21 -0400 Subject: [PATCH 109/120] up --- src/metrics.jl | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/metrics.jl b/src/metrics.jl index 75f2dfc..615d1bf 100644 --- a/src/metrics.jl +++ b/src/metrics.jl @@ -160,22 +160,22 @@ function correlation(m, x, y, w; agg=mean) p = view(m(x), 1, :) y = vec(y) w = vec(w) - p_mean = w' * p - p_var = w' * (p .^ 2) - p_mean^2 - y_mean = w' * y - y_var = w' * (y .^ 2) - y_mean^2 - py_mean = w' * (p .* y) + p_mean = w' * p / sum(w) + p_var = w' * (p .^ 2) / sum(w) - p_mean^2 + y_mean = w' * y / sum(w) + y_var = w' * (y .^ 2) / sum(w) - y_mean^2 + py_mean = w' * (p .* y) / sum(w) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) end function correlation(m, x, y, w, offset; agg=mean) p = view(m(x), 1, :) .+ view(offset, 1, :) y = vec(y) w = vec(w) - p_mean = w' * p - p_var = w' * (p .^ 2) - p_mean^2 - y_mean = w' * y - y_var = w' * (y .^ 2) - y_mean^2 - py_mean = w' * (p .* y) + p_mean = w' * p / sum(w) + p_var = w' * (p .^ 2) / sum(w) - p_mean^2 + y_mean = w' * y / sum(w) + y_var = w' * (y .^ 2) / sum(w) - y_mean^2 + py_mean = w' * (p .* y) / sum(w) return (py_mean - p_mean * y_mean) / (sqrt(p_var) * sqrt(y_var)) * sum(w) end From af6ae4cadddfd8ebb7706b8a7aba38462d149ccb Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Mon, 4 May 2026 09:45:56 +0530 Subject: [PATCH 110/120] rewrite mlogloss with ifelse (#45) Signed-off-by: AdityaPandeyCN --- src/losses.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/losses.jl b/src/losses.jl index b3a0ef6..bb9f5f4 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -47,8 +47,8 @@ function _mlogloss_core(pred, y) nclasses = size(pred, 1) classes = reshape(Int32(1):Int32(nclasses), :, 1, 1) y_idx = reshape(Int32.(y), 1, 1, :) - y_oh = Float32.(classes .== y_idx) - return -sum(y_oh .* logsoftmax(pred; dims=1); dims=1) + lsm = logsoftmax(pred; dims=1) + return -sum(ifelse.(classes .== y_idx, lsm, zero(eltype(lsm))); dims=1) end function _tweedie_core(pred, y) From d9cf82bfcb8338464bd9ef8db53cbf67a033c4c4 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Mon, 4 May 2026 00:16:36 -0400 Subject: [PATCH 111/120] up --- benchmarks/titanic-logloss.jl | 10 +++++----- benchmarks/titanic-mlogloss.jl | 18 +++++++++++++----- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index c6d8cf1..53ece98 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -38,11 +38,11 @@ feature_names = setdiff(names(df), ["Survived"]) arch = NeuroTabModels.NeuroTreeConfig(; tree_type=:binary, - k=8, + k=1, depth=4, ntrees=16, - stack_size=2, - hidden_size=8, + stack_size=1, + hidden_size=1, actA=:identity, init_scale=1.0, scaler=true, @@ -82,8 +82,8 @@ learner = NeuroTabRegressor( nrounds=200, early_stopping_rounds=2, lr=1e-2, - device=:reactant, - backend=:enzyme + device=:cpu, + backend=:reactant ) @time m = NeuroTabModels.fit( diff --git a/benchmarks/titanic-mlogloss.jl b/benchmarks/titanic-mlogloss.jl index e3cae01..3d9ef23 100644 --- a/benchmarks/titanic-mlogloss.jl +++ b/benchmarks/titanic-mlogloss.jl @@ -7,13 +7,15 @@ using CategoricalArrays using Random using CategoricalArrays +using Enzyme +using Reactant +# using CUDA, cuDNN +using Zygote + Random.seed!(123) df = MLDatasets.Titanic().dataframe -# convert target variable to a categorical -transform!(df, :Survived => categorical => :y_cat) - # convert string feature to Categorical transform!(df, :Sex => categorical => :Sex) transform!(df, :Sex => ByRow(levelcode) => :Sex) @@ -25,6 +27,9 @@ transform!(df, :Age => (x -> coalesce.(x, median(skipmissing(x)))) => :Age); # remove unneeded variables df = df[:, Not([:PassengerId, :Name, :Embarked, :Cabin, :Ticket])] +# convert target variable to a categorical +transform!(df, :Survived => categorical => :y_cat) + train_ratio = 0.8 train_indices = randperm(nrow(df))[1:Int(round(train_ratio * nrow(df)))] @@ -38,10 +43,12 @@ eltype(dtrain[:, "y_cat"]) arch = NeuroTabModels.NeuroTreeConfig(; actA=:identity, init_scale=1.0, + k=8, depth=4, ntrees=16, stack_size=1, hidden_size=1, + scaler=true ) # arch = NeuroTabModels.TabMConfig(; @@ -59,10 +66,11 @@ arch = NeuroTabModels.NeuroTreeConfig(; learner = NeuroTabClassifier( arch; - nrounds=100, + nrounds=200, early_stopping_rounds=2, - lr=1e-2, + lr=2e-2, device=:cpu, + backend=:reactant ) m = NeuroTabModels.fit( From c87095c4c022073ee5b747af39544870a4057085 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Mon, 8 Jun 2026 21:29:01 -0400 Subject: [PATCH 112/120] up --- src/models/MLP/mlp.jl | 4 ++-- src/models/ResNet/resnet.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/models/MLP/mlp.jl b/src/models/MLP/mlp.jl index b61b781..fb7bb0f 100644 --- a/src/models/MLP/mlp.jl +++ b/src/models/MLP/mlp.jl @@ -4,8 +4,8 @@ export MLPConfig import Flux import Flux: @functor, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: logsoftmax, softmax, softmax!, relu, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch -import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel, SkipConnection +import Flux: relu, gelu, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold +import Flux: BatchNorm, Dense, Dropout, Parallel, SkipConnection import ..Models: get_loss_type, GaussianMLE import ..Models: Architecture diff --git a/src/models/ResNet/resnet.jl b/src/models/ResNet/resnet.jl index 6a6a263..92f5702 100644 --- a/src/models/ResNet/resnet.jl +++ b/src/models/ResNet/resnet.jl @@ -4,8 +4,8 @@ export ResNetConfig import Flux import Flux: @functor, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: logsoftmax, softmax, softmax!, relu, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold, onehotbatch -import Flux: BatchNorm, Dense, Dropout, MultiHeadAttention, Parallel, SkipConnection +import Flux: relu, gelu, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold +import Flux: BatchNorm, Dense, Dropout, Parallel, SkipConnection import ..Models: get_loss_type, GaussianMLE import ..Models: Architecture From d436d472be2eca28cff9953098429ce7e6d226c7 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 9 Jun 2026 01:39:11 -0400 Subject: [PATCH 113/120] up --- Project.toml | 2 +- benchmarks/Project.toml | 4 +- benchmarks/YEAR-regression-grp.jl | 2 +- benchmarks/YEAR-regression.jl | 20 +++-- src/models/MLP/mlp.jl | 80 +++++++---------- src/models/ResNet/resnet.jl | 142 ++++++++++-------------------- src/models/models.jl | 21 ++++- test/core.jl | 42 +++++++++ 8 files changed, 153 insertions(+), 160 deletions(-) diff --git a/Project.toml b/Project.toml index 61e6f2e..518c54a 100644 --- a/Project.toml +++ b/Project.toml @@ -32,7 +32,7 @@ NNlib = "0.9" OneHotArrays = "0.2" Optimisers = "0.4" Random = "1" -Reactant = "0.2.254" +Reactant = "0.2.264" Statistics = "1" StatsBase = "0.34" Tables = "1.9" diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 6ab1d45..d0f4266 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -16,5 +16,5 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" [compat] -CUDA = "5.11" -cuDNN = "1" \ No newline at end of file +CUDA = "6.1" +cuDNN = "6.1" \ No newline at end of file diff --git a/benchmarks/YEAR-regression-grp.jl b/benchmarks/YEAR-regression-grp.jl index f44228e..74e9366 100644 --- a/benchmarks/YEAR-regression-grp.jl +++ b/benchmarks/YEAR-regression-grp.jl @@ -83,7 +83,7 @@ arch = NeuroTabModels.NeuroTreeConfig(; # hidden_size=256, # ) # arch = NeuroTabModels.ResNetConfig(; -# num_blocks=1, +# stack_size=1, # hidden_size=128, # act=:relu, # dropout=0.5, diff --git a/benchmarks/YEAR-regression.jl b/benchmarks/YEAR-regression.jl index 569a876..fa800a1 100644 --- a/benchmarks/YEAR-regression.jl +++ b/benchmarks/YEAR-regression.jl @@ -76,19 +76,21 @@ arch = NeuroTabModels.NeuroTreeConfig(; # dropout=0.1, # # scaling_init=:normal, # ) + # arch = NeuroTabModels.MLPConfig(; # act=:relu, -# stack_size=1, -# hidden_size=256, -# ) -# arch = NeuroTabModels.ResNetConfig(; -# num_blocks=1, +# stack_size=2, # hidden_size=128, -# act=:relu, -# dropout=0.5, -# MLE_tree_split=false +# dropout=0.2 # ) +arch = NeuroTabModels.ResNetConfig(; + stack_size=2, + hidden_size=128, + act=:relu, + dropout=0.5, +) + device = :gpu backend = :reactant loss = :mse # :mse :gaussian_mle :tweedie @@ -111,7 +113,7 @@ learner = NeuroTabRegressor( nrounds=200, early_stopping_rounds=2, lr=1e-3, - batchsize=512, + batchsize=1024, device, backend ) diff --git a/src/models/MLP/mlp.jl b/src/models/MLP/mlp.jl index fb7bb0f..380f12f 100644 --- a/src/models/MLP/mlp.jl +++ b/src/models/MLP/mlp.jl @@ -2,93 +2,81 @@ module MLP export MLPConfig -import Flux -import Flux: @functor, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: relu, gelu, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold -import Flux: BatchNorm, Dense, Dropout, Parallel, SkipConnection +using Lux +using LuxCore -import ..Models: get_loss_type, GaussianMLE -import ..Models: Architecture +import ..Models: Architecture, get_activation struct MLPConfig <: Architecture act::Symbol hidden_size::Int stack_size::Int + dropout::Float64 MLE_tree_split::Bool end function MLPConfig(; kwargs...) - - # defaults arguments args = Dict{Symbol,Any}( :act => :relu, :hidden_size => 64, :stack_size => 1, - :MLE_tree_split => false + :dropout => 0.0, + :MLE_tree_split => false, ) args_ignored = setdiff(keys(kwargs), keys(args)) - args_ignored_str = join(args_ignored, ", ") length(args_ignored) > 0 && - @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(args_ignored_str)." + @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(join(args_ignored, ", "))." args_default = setdiff(keys(args), keys(kwargs)) - args_default_str = join(args_default, ", ") length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(args_default_str)." + @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." - args_override = intersect(keys(args), keys(kwargs)) - for arg in args_override + for arg in intersect(keys(args), keys(kwargs)) args[arg] = kwargs[arg] end - config = MLPConfig( + return MLPConfig( Symbol(args[:act]), args[:hidden_size], args[:stack_size], - args[:MLE_tree_split] + args[:dropout], + args[:MLE_tree_split], ) - - return config end -function (config::MLPConfig)(; nfeats, outsize) +function _mlp_trunk(nfeats::Int, hsize::Int, outsize::Int, act, stack_size::Int, dropout::Float64) + layers = Any[ + Dense(nfeats => hsize), + ] + for _ in 1:stack_size + push!(layers, BatchNorm(hsize, act)) + push!(layers, Dense(hsize => hsize)) + dropout > 0 && push!(layers, Dropout(dropout)) + end + push!(layers, Dense(hsize => outsize)) + return Chain(layers...) +end +function (config::MLPConfig)(; nfeats, outsize, kwargs...) + act = get_activation(config.act) hsize = config.hidden_size - if config.MLE_tree_split && outsize == 2 - outsize ÷= 2 + if config.MLE_tree_split + iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") + head_outsize = outsize ÷ 2 chain = Chain( - BatchNorm(nfeats), - Dense(nfeats => hsize), Parallel( vcat, - Chain( - BatchNorm(nfeats), - Dense(nfeats => hsize), - BatchNorm(hsize, relu), - Dense(hsize => outsize) - ), - Chain( - BatchNorm(nfeats), - Dense(nfeats => hsize), - BatchNorm(hsize, relu), - Dense(hsize => outsize) - ) - ) + _mlp_trunk(nfeats, hsize, head_outsize, act, config.stack_size, config.dropout), + _mlp_trunk(nfeats, hsize, head_outsize, act, config.stack_size, config.dropout), + ), ) else - chain = Chain( - BatchNorm(nfeats), - Dense(nfeats => hsize), - BatchNorm(hsize, relu), - Dense(hsize => hsize), - BatchNorm(hsize, relu), - Dense(hsize => outsize) - ) + chain = _mlp_trunk(nfeats, hsize, outsize, act, config.stack_size, config.dropout) end return chain end -end \ No newline at end of file +end diff --git a/src/models/ResNet/resnet.jl b/src/models/ResNet/resnet.jl index 92f5702..66749ea 100644 --- a/src/models/ResNet/resnet.jl +++ b/src/models/ResNet/resnet.jl @@ -2,16 +2,13 @@ module ResNet export ResNetConfig -import Flux -import Flux: @functor, trainmode!, gradient, Chain, DataLoader, cpu, gpu -import Flux: relu, gelu, sigmoid, sigmoid_fast, hardsigmoid, tanh, tanh_fast, hardtanh, softplus, onecold -import Flux: BatchNorm, Dense, Dropout, Parallel, SkipConnection +using Lux +using LuxCore -import ..Models: get_loss_type, GaussianMLE -import ..Models: Architecture +import ..Models: Architecture, get_activation struct ResNetConfig <: Architecture - num_blocks::Int + stack_size::Int hidden_size::Int act::Symbol dropout::Float64 @@ -19,130 +16,81 @@ struct ResNetConfig <: Architecture end function ResNetConfig(; kwargs...) - - # defaults arguments args = Dict{Symbol,Any}( - :num_blocks => 1, + :stack_size => 1, :hidden_size => 64, :act => :relu, - :dropout => 1.0, - :MLE_tree_split => false + :dropout => 0.0, + :MLE_tree_split => false, ) args_ignored = setdiff(keys(kwargs), keys(args)) - args_ignored_str = join(args_ignored, ", ") length(args_ignored) > 0 && - @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(args_ignored_str)." + @warn "Following $(length(args_ignored)) provided arguments will be ignored: $(join(args_ignored, ", "))." args_default = setdiff(keys(args), keys(kwargs)) - args_default_str = join(args_default, ", ") length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(args_default_str)." + @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." - args_override = intersect(keys(args), keys(kwargs)) - for arg in args_override + for arg in intersect(keys(args), keys(kwargs)) args[arg] = kwargs[arg] end - config = ResNetConfig( - args[:num_blocks], + return ResNetConfig( + args[:stack_size], args[:hidden_size], Symbol(args[:act]), args[:dropout], - args[:MLE_tree_split] + args[:MLE_tree_split], ) - - return config end -function ResBlock_v1(; hsize, dropout) - layer = SkipConnection( - Chain( - BatchNorm(hsize), - Dense(hsize => hsize, relu), - Dropout(dropout), - Dense(hsize => hsize), - Dropout(dropout), - ), - + - ) - return layer -end -function ResBlock_v2A(; hsize, dropout, kwargs...) - layer = Chain( - SkipConnection( - Chain( - Dense(hsize => hsize), - BatchNorm(hsize, relu), - Dropout(dropout), - Dense(hsize => hsize), - BatchNorm(hsize), - ), - + - ), - x -> relu.(x) - ) - return layer -end -function ResBlock_v2B(; hsize, dropout, kwargs...) - layer = Chain( - SkipConnection( - Chain( - Dense(hsize => hsize), - BatchNorm(hsize, relu), - Dropout(dropout), - Dense(hsize => hsize), - BatchNorm(hsize), - ), - vcat - ), - Dense(2 * hsize => hsize, relu), +function _res_block(hsize::Int, act, dropout::Float64) + layers = Any[ + Dense(hsize => hsize), + BatchNorm(hsize, act), + ] + dropout > 0 && push!(layers, Dropout(dropout)) + push!(layers, Dense(hsize => hsize)) + push!(layers, BatchNorm(hsize)) + + return Chain( + SkipConnection(Chain(layers...), +), + BatchNorm(hsize, act), ) - return layer end +function _resnet_trunk(nfeats::Int, hsize::Int, outsize::Int, act, stack_size::Int, dropout::Float64) + layers = Any[ + Dense(nfeats => hsize), + BatchNorm(hsize, act), + ] + for _ in 1:stack_size + push!(layers, _res_block(hsize, act, dropout)) + end + push!(layers, Dense(hsize => outsize)) + return Chain(layers...) +end -function (config::ResNetConfig)(; nfeats, outsize) - +function (config::ResNetConfig)(; nfeats, outsize, kwargs...) + act = get_activation(config.act) hsize = config.hidden_size - dropout = config.dropout - if config.MLE_tree_split && outsize == 2 - outsize ÷= 2 + if config.MLE_tree_split + iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") + head_outsize = outsize ÷ 2 chain = Chain( - BatchNorm(nfeats), - Dense(nfeats => hsize), - BatchNorm(hsize, relu), Parallel( vcat, - Chain( - ResBlock_v2A(; hsize, dropout), - # BatchNorm(hsize), - Dense(hsize => outsize), - # BatchNorm(outsize), - ), - Chain( - ResBlock_v2A(; hsize, dropout), - # BatchNorm(hsize), - Dense(hsize => outsize), - # BatchNorm(outsize), - ) + _resnet_trunk(nfeats, hsize, head_outsize, act, config.stack_size, config.dropout), + _resnet_trunk(nfeats, hsize, head_outsize, act, config.stack_size, config.dropout), ), ) else - chain = Chain( - BatchNorm(nfeats), - Dense(nfeats => hsize), - BatchNorm(hsize, relu), - ResBlock_v2A(; hsize, dropout), - # BatchNorm(hsize), - Dense(hsize => outsize), - # BatchNorm(outsize), - ) - @info "single chain" + chain = _resnet_trunk(nfeats, hsize, outsize, act, config.stack_size, config.dropout) end return chain end -end \ No newline at end of file +end diff --git a/src/models/models.jl b/src/models/models.jl index ad8b5cd..a1f3725 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -12,6 +12,19 @@ abstract type Architecture end _broadcast_relu(x) = NNlib.relu.(x) +const activation_dict = Dict{Symbol,Function}( + :relu => NNlib.relu, + :gelu => NNlib.gelu, + :sigmoid => NNlib.sigmoid_fast, + :tanh => NNlib.tanh_fast, +) + +function get_activation(act::Symbol) + haskey(activation_dict, act) || + error("Unknown activation: $act. Supported: $(sort(collect(keys(activation_dict))))") + return activation_dict[act] +end + """ NeuroTabModel @@ -41,10 +54,10 @@ using .MOETrees include("TabM/TabM.jl") using .TabM -# include("MLP/mlp.jl") -# using .MLP +include("MLP/mlp.jl") +using .MLP -# include("ResNet/resnet.jl") -# using .ResNet +include("ResNet/resnet.jl") +using .ResNet end \ No newline at end of file diff --git a/test/core.jl b/test/core.jl index b89d29c..c9eb03f 100644 --- a/test/core.jl +++ b/test/core.jl @@ -210,6 +210,48 @@ end end +@testset "Classification - $arch_name" for (arch_name, arch) in [ + ("MLP", NeuroTabModels.MLPConfig(; hidden_size=32, stack_size=1)), + ("ResNet", NeuroTabModels.ResNetConfig(; hidden_size=32, stack_size=1)), +] + + Random.seed!(123) + X, y = @load_crabs + df = DataFrame(X) + df[!, :class] = y + target_name = "class" + feature_names = setdiff(names(df), [target_name]) + + train_ratio = 0.8 + train_indices = randperm(nrow(df))[1:Int(train_ratio * nrow(df))] + + dtrain = df[train_indices, :] + deval = df[setdiff(1:nrow(df), train_indices), :] + + learner = NeuroTabClassifier(arch; + embedding_config=Dict(:embedding_type => :batchnorm), + nrounds=200, + batchsize=32, + early_stopping_rounds=5, + lr=1e-2, + ) + + m = NeuroTabModels.fit( + learner, + dtrain; + deval, + target_name, + feature_names, + print_every_n=5 + ) + + ptrain = [argmax(x) for x in eachrow(m(dtrain))] + peval = [argmax(x) for x in eachrow(m(deval))] + @test mean(ptrain .== levelcode.(dtrain.class)) > 0.95 + @test mean(peval .== levelcode.(deval.class)) > 0.95 + +end + @testset "Backend/device - reactant is a backend" begin @test_throws ErrorException NeuroTabRegressor(; backend=:zygote, device=:reactant) @test_throws ErrorException NeuroTabClassifier(; backend=:zygote, device=:reactant) From afa5857924832cb0195b9ad874576011d3ab0801 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 9 Jun 2026 22:49:18 -0400 Subject: [PATCH 114/120] up --- docs/make.jl | 2 + docs/package.json | 20 +- docs/src/.vitepress/theme/style.css | 343 ++++++++++++++++++---------- docs/src/embeddings.md | 5 + docs/src/models/mlp.md | 5 + docs/src/models/neurotrees.md | 4 +- docs/src/models/resnet.md | 6 + docs/src/models/tabM.md | 4 +- 8 files changed, 256 insertions(+), 133 deletions(-) create mode 100644 docs/src/embeddings.md create mode 100644 docs/src/models/mlp.md create mode 100644 docs/src/models/resnet.md diff --git a/docs/make.jl b/docs/make.jl index 8670f59..0b86a9f 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -8,6 +8,8 @@ pages = [ "Design" => "design.md", "Models" => [ "Interface" => "models/models.md", + "MLP" => "models/mlp.md", + "ResNet" => "models/resnet.md", "NeuroTrees" => "models/neurotrees.md", "TabM" => "models/tabM.md", ], diff --git a/docs/package.json b/docs/package.json index 3a29a27..cfca9a4 100644 --- a/docs/package.json +++ b/docs/package.json @@ -1,12 +1,7 @@ { "devDependencies": { - "@nolebase/vitepress-plugin-enhanced-readabilities": "^2.15.0", - "@types/d3-format": "^3.0.4", - "@types/node": "^22.13.9", - "markdown-it": "^14.1.0", - "markdown-it-mathjax3": "^4.3.2", - "vitepress": "^1.6.3", - "vitepress-plugin-tabs": "^0.6.0" + "@types/markdown-it-footnote": "^3.0.4", + "@types/node": "^25.3.5" }, "scripts": { "docs:dev": "vitepress dev build/.documenter", @@ -14,7 +9,12 @@ "docs:preview": "vitepress preview build/.documenter" }, "dependencies": { - "d3-format": "^3.1.0", - "markdown-it-footnote": "^4.0.0" + "@mdit/plugin-mathjax": "^0.26.1", + "@mdit/plugin-tex": "^0.24.1", + "@nolebase/vitepress-plugin-enhanced-readabilities": "^2.18.2", + "markdown-it": "^14.1.0", + "markdown-it-footnote": "^4.0.0", + "vitepress": "^1.6.4", + "vitepress-plugin-tabs": "^0.9.0" } -} \ No newline at end of file +} diff --git a/docs/src/.vitepress/theme/style.css b/docs/src/.vitepress/theme/style.css index fb16117..2fa6fe7 100644 --- a/docs/src/.vitepress/theme/style.css +++ b/docs/src/.vitepress/theme/style.css @@ -1,16 +1,21 @@ +/* Customize default theme styling by overriding CSS variables: +https://github.com/vuejs/vitepress/blob/main/src/client/theme-default/styles/vars.css */ +/* Example */ +/* https://github.com/vuejs/vitepress/blob/main/template/.vitepress/theme/style.css */ + .VPHero .clip { white-space: pre; max-width: 600px; } /* Fonts */ - @font-face { - font-family: JuliaMono-Regular; - src: url("https://cdn.jsdelivr.net/gh/cormullion/juliamono/webfonts/JuliaMono-Regular.woff2"); + font-family: JuliaMono-Regular; + src: url("https://cdn.jsdelivr.net/gh/cormullion/juliamono/webfonts/JuliaMono-Regular.woff2"); } :root { + scroll-behavior: smooth; /* Typography */ --vp-font-family-base: "Barlow", "Inter var experimental", "Inter var", -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Oxygen, Ubuntu, @@ -21,143 +26,122 @@ } /* Disable contextual alternates (kind of like ligatures but different) in monospace, - which turns `/>` to an up arrow and `|>` (the Julia pipe symbol) to an up arrow as well. */ + which turns `/>` to an up arrow and `|>` (the Julia pipe symbol) to an up arrow as well. */ .mono-no-substitutions { - font-family: "JuliaMono-Light", monospace; + font-family: "JuliaMono-Regular", monospace; font-feature-settings: "calt" off; } .mono-no-substitutions-alt { - font-family: "JuliaMono-Light", monospace; + font-family: "JuliaMono-Regular", monospace; font-variant-ligatures: none; } -pre, code { - font-family: "JuliaMono-Light", monospace; +pre, +code { + font-family: "JuliaMono-Regular", monospace; font-feature-settings: "calt" off; } /* Colors */ - :root { --julia-blue: #4063D8; --julia-purple: #9558B2; --julia-red: #CB3C33; --julia-green: #389826; - --vp-c-brand: #40a52f; - --vp-c-brand-light: #3dd027; - --vp-c-brand-lighter: #9499ff; - --vp-c-brand-lightest: #bcc0ff; - --vp-c-brand-dark: #535bf2; - --vp-c-brand-darker: #454ce1; + --vp-c-brand: #0087d7; + --vp-c-brand-1: #0890df; + --vp-c-brand-2: #0599ef; + --vp-c-brand-3: #0c9ff4; + --vp-c-brand-light: #0087d7; + --vp-c-brand-dark: #5fd7ff; --vp-c-brand-dimm: #212425; - --vp-tip-bg: rgb(254, 254, 254); + + --vp-evo-blue: #023572; + --vp-evo-gray: #c9ccd1; + --vp-evo-gray-dark: #454e56; + --vp-evo-blue-light: #5891d5; + --vp-evo-blue-dark: #4571a5; + --vp-evo-green: #26a671; + --vp-evo-green-hover: #1e855a ; + --vp-evo-red: #e5616c; + --vp-evo-yellow: #ffc700; + --vp-evo-backgorund-alt: #f9f9f9; + --vp-evo-backgorund-light: #f0f3f7; + --vp-evo-backgorund-medium: #e6ebf1; + --vp-evo-backgorund-dark: #c9ccd1; + /* Greens */ - --vp-dark-green: #155f3e; /* Main accent green */ - --vp-dark-green-dark: #26a671; + --vp-dark-green: #155f3e; + /* Main accent green */ + --vp-dark-green-dark: #2b855c; --vp-dark-green-light: #42d392; --vp-dark-green-lighter: #35eb9a; - --vp-dark-green-dimm: #2e3834; - /* Complementary Colors */ --vp-dark-gray: #1e1e1e; --vp-dark-gray-soft: #2a2a2a; --vp-dark-gray-mute: #242424; --vp-light-gray: #d1d5db; + --vp-tip-bg: rgb(254, 254, 254); /* Text Colors */ - --vp-dark-text: #e5e5e5; /* Primary text color */ - --vp-dark-subtext: #c1c1c1; /* Subtle text */ - - /* Backgrounds */ - --vp-dark-bg: #121212; /* Main background */ - --vp-dark-bg-alt: #1a1a1a; /* Alternative background */ - --vp-dark-bg-dimm: #101010; - - /* Borders */ - --vp-dark-border: #3b8a65; - /* custom tip */ - --vp-custom-block-tip-border: var(--vp-c-brand-darker); + --vp-dark-text: #e5e5e5; + /* Primary text color */ + --vp-dark-subtext: #c1c1c1; + /* Subtle text */ + --vp-source-text: #e5e5e5; + /* custom tip */ + --vp-custom-block-tip-border: var(--vp-c-brand-light); --vp-custom-block-tip-bg: var(--vp-tip-bg); - /* button */ - --vp-button-brand-border: var(--vp-dark-green-dark); - --vp-button-brand-bg: var(--vp-dark-green-dark); - --vp-button-brand-hover-border: var(--vp-dark-green-dark); - --vp-button-brand-hover-bg: var(--vp-dark-green); - --vp-button-brand-active-border: var(--vp-dark-green-light); - --vp-button-brand-active-bg: var(--vp-dark-green); } +/* Component: Button */ +:root { + --vp-button-brand-border: var(--vp-evo-green); + --vp-button-brand-bg: var(--vp-evo-green); + --vp-button-brand-hover-border: var(--vp-evo-blue-light); + --vp-button-brand-hover-bg: var(--vp-evo-blue-light); + --vp-button-brand-active-border: var(--vp-evo-blue-light); + --vp-button-brand-active-bg: var(--vp-evo-blue-light); + + --vp-button-alt-border: var(--vp-evo-gray-dark); + --vp-button-alt-bg: var(--vp-evo-gray-dark); + --vp-button-alt-text: var(--vp-c-white); + --vp-button-alt-hover-border: var(--vp-evo-blue-light); + --vp-button-alt-hover-bg: var(--vp-evo-blue-light); + --vp-button-alt-hover-text: var(--vp-c-white); + --vp-button-alt-active-border: var(--vp-evo-blue-light); + --vp-button-alt-active-bg: var(--vp-evo-blue-light); + --vp-button-alt-active-text: var(--vp-c-white); +} + +/* Component: Home */ +:root { + --vp-home-hero-name-color: var(--vp-evo-blue); +} + +/* Hero Section & dark theme overrides */ :root.dark { + --vp-home-hero-name-color: var(--vp-evo-blue-light); + /* custom tip */ --vp-custom-block-tip-border: var(--vp-dark-green-dark); --vp-custom-block-tip-text: var(--vp-dark-subtext); --vp-custom-block-tip-bg: var(--vp-dark-gray-mute); } -/* Button */ -:root.dark { - --vp-button-brand-border: var(--vp-dark-green-light); - --vp-button-brand-text: var(--vp-dark-text); - --vp-button-brand-bg: var(--vp-dark-green); - --vp-button-brand-hover-border: var(--vp-dark-green-lighter); - --vp-button-brand-hover-text: var(--vp-dark-text); - --vp-button-brand-hover-bg: var(--vp-dark-green-dark); - --vp-button-brand-active-border: var(--vp-dark-green-light); - --vp-button-brand-active-text: var(--vp-dark-text); - --vp-button-brand-active-bg: var(--vp-dark-green); -} - -/* Typography */ -:root.dark { - --vp-dark-link: var(--vp-dark-green-lighter); - --vp-dark-link-hover: var(--vp-dark-green-light); -} - -/* Hero Section */ -/* :root.dark { - --vp-home-hero-name-color: transparent; - --vp-home-hero-name-background: -webkit-linear-gradient( - 120deg, - var(--julia-purple) 15%, - var(--vp-dark-green-light), - var(--vp-dark-green) - ); - --vp-home-hero-image-background-image: linear-gradient( - -45deg, - var(--vp-dark-green) 30%, - var(--vp-dark-green-light), - var(--vp-dark-gray) 30% - ); - --vp-home-hero-image-filter: blur(56px); -} */ +/** + * Colors links + * -------------------------------------------------------------------------- */ .dark { + --vp-c-brand: var(--vp-dark-green-light); --vp-c-brand-1: var(--vp-dark-green-light); --vp-c-brand-2: var(--vp-dark-green-lighter); --vp-c-brand-3: var(--vp-dark-green); - } -/* Component: Home */ -/* :root { - --vp-home-hero-name-color: transparent; - --vp-home-hero-name-background: -webkit-linear-gradient( - 120deg, - #9558B2 30%, - #CB3C33 - ); - - --vp-home-hero-image-background-image: linear-gradient( - -45deg, - #9558b282 30%, - #3798269a 30%, - #cb3d33e3 - ); - --vp-home-hero-image-filter: blur(40px); -} */ - @media (min-width: 640px) { :root { --vp-home-hero-image-filter: blur(56px); @@ -170,46 +154,167 @@ pre, code { } } -/* Component: Algolia */ - -.DocSearch { - --docsearch-primary-color: var(--vp-c-brand) !important; -} - /* Component: MathJax */ -mjx-container > svg { + +.mjx-scroll-wrapper { display: block; - margin: auto; + max-width: 100%; + overflow-x: auto; + overflow-y: hidden; + text-align: center; } mjx-container { + display: inline-block; padding: 0.5rem 0; -} - -mjx-container { - display: inline; margin: auto 2px -2px; + overflow: visible !important; + /* Override MathJax's overflow */ + text-align: left; } -mjx-container > svg { - margin: auto; +mjx-container>svg { display: inline-block; + height: auto; } -/* Component: Docstring Custom Block */ +mjx-container>svg[data-labels="true"] { + position: absolute; + left: 0; + top: 0; + width: 100%; + height: 100%; + pointer-events: none; +} + +/* @mdit/plugin-mathjax theming refs */ -.jldocstring.custom-block { - border: 1px solid var(--vp-c-gray-2); - color: var(--vp-c-text-1); +mjx-container svg path { + fill: currentColor !important; + stroke: currentColor !important; } -.jldocstring.custom-block summary { - font-weight: 700; - cursor: pointer; - user-select: none; - margin: 0 0 8px; +.MathJax_ref { + color: var(--vp-c-brand-1); } -.jldocstring.custom-block summary a { - pointer-events: none; + +.MathJax_ref:hover { + color: var(--vp-c-brand-2); + cursor: pointer; +} + +/* ===== Custom colors for input and output code blocks ===== */ +:root { + /* --vp-c-bg-input: #eef0f3; */ + --vp-c-bg-output: #fbfbfb; + --vp-c-bg-output-outline: #e2e2e3; +} + +.dark { + /* --vp-c-bg-input: #1a1a1a; */ + --vp-c-bg-output: #1a1a1a; + --vp-c-bg-output-outline: #2e2e32; +} + +/* +.language-julia { + background-color: var(--vp-c-bg-input) !important; +} +*/ + +.language- { + background-color: var(--vp-c-bg-output) !important; + outline: 1px solid var(--vp-c-bg-output-outline); +} + +/* Julia REPL prompt syntax highlighting */ +/* prompt token behavior */ +.vp-doc .shiki .repl-prompt { + opacity: 1.0; + user-select: none; +} + +:root:not(.dark) .vp-doc .shiki .repl-prompt-julia { + color: var(--vp-dark-green); } + +:root:not(.dark) .vp-doc .shiki .repl-prompt-pkg { + color: var(--vp-c-brand-light); +} + +.dark .vp-doc .shiki .repl-prompt-julia { + color: var(--vp-dark-green-light); +} + +.dark .vp-doc .shiki .repl-prompt-pkg { + color: var(--vp-c-brand-dark); +} + +/* ===== Sidebar Drawer Toggle ===== */ +/* Smooth transitions for sidebar collapse/expand on desktop */ +@media (min-width: 960px) { + .VPSidebar { + transition: transform 0.35s cubic-bezier(0.4, 0, 0.2, 1), + opacity 0.25s ease; + } + + .VPContent.has-sidebar { + transition: padding-left 0.35s cubic-bezier(0.4, 0, 0.2, 1), + padding-right 0.35s cubic-bezier(0.4, 0, 0.2, 1); + } + + /* Smooth transitions for navbar elements (needed for expand-back too) */ + .VPNavBar.has-sidebar .content { + transition: padding-left 0.35s cubic-bezier(0.4, 0, 0.2, 1) !important; + } + + .VPNavBar.has-sidebar .divider { + transition: padding-left 0.35s cubic-bezier(0.4, 0, 0.2, 1) !important; + } + + .VPNavBar.has-sidebar .title { + transition: width 0.35s cubic-bezier(0.4, 0, 0.2, 1), + padding 0.35s cubic-bezier(0.4, 0, 0.2, 1), + opacity 0.25s ease !important; + } + + /* Collapsed state: zero the sidebar width variable so all layout formulas + (VitePress defaults + Enhanced Readabilities plugin) adjust automatically */ + html.sidebar-drawer-collapsed { + --vp-sidebar-width: 0px !important; + } + + html.sidebar-drawer-collapsed .VPSidebar { + transform: translateX(-100%); + opacity: 0; + pointer-events: none; + } + + /* Original width mode only: let content fill space up to the aside */ + html.sidebar-drawer-collapsed body:not(.VPNolebaseEnhancedReadabilitiesLayoutSwitchFullWidth):not(.VPNolebaseEnhancedReadabilitiesLayoutSwitchSidebarWidthAdjustableOnly):not(.VPNolebaseEnhancedReadabilitiesLayoutSwitchBothWidthAdjustable) .VPDoc.has-sidebar .container { + display: flex; + justify-content: center; + max-width: 992px; + } + + html.sidebar-drawer-collapsed body:not(.VPNolebaseEnhancedReadabilitiesLayoutSwitchFullWidth):not(.VPNolebaseEnhancedReadabilitiesLayoutSwitchSidebarWidthAdjustableOnly):not(.VPNolebaseEnhancedReadabilitiesLayoutSwitchBothWidthAdjustable) .VPDoc.has-sidebar .content-container { + max-width: none; + } + + /* Collapse the navbar title area (site logo + name) */ + html.sidebar-drawer-collapsed .VPNavBar.has-sidebar .title, + html.sidebar-drawer-collapsed .VPNavBar.has-sidebar>.wrapper>.container>.title { + width: 0 !important; + overflow: hidden; + opacity: 0; + padding: 0 !important; + pointer-events: none; + } + + /* Move drawer button closer to the edge */ + html.sidebar-drawer-collapsed .VPNavBar.has-sidebar .content, + html.sidebar-drawer-collapsed .VPNavBar.has-sidebar>.wrapper>.container>.content { + padding-left: 8px !important; + } +} \ No newline at end of file diff --git a/docs/src/embeddings.md b/docs/src/embeddings.md new file mode 100644 index 0000000..5d4d933 --- /dev/null +++ b/docs/src/embeddings.md @@ -0,0 +1,5 @@ +# Embeddings + +```@autodocs +Modules = [NeuroTabModels.Models.Embeddings] +``` diff --git a/docs/src/models/mlp.md b/docs/src/models/mlp.md new file mode 100644 index 0000000..7f8a172 --- /dev/null +++ b/docs/src/models/mlp.md @@ -0,0 +1,5 @@ +# MLP + +```@autodocs +Modules = [NeuroTabModels.Models.MLP] +``` \ No newline at end of file diff --git a/docs/src/models/neurotrees.md b/docs/src/models/neurotrees.md index 1787c2a..34883d8 100644 --- a/docs/src/models/neurotrees.md +++ b/docs/src/models/neurotrees.md @@ -1,5 +1,5 @@ # NeuroTrees -```@docs -NeuroTreeConfig +```@autodocs +Modules = [NeuroTabModels.Models.NeuroTrees] ``` diff --git a/docs/src/models/resnet.md b/docs/src/models/resnet.md new file mode 100644 index 0000000..bf911e9 --- /dev/null +++ b/docs/src/models/resnet.md @@ -0,0 +1,6 @@ +# ResNet + + +```@autodocs +Modules = [NeuroTabModels.Models.ResNet] +``` \ No newline at end of file diff --git a/docs/src/models/tabM.md b/docs/src/models/tabM.md index e3300b6..e0e3385 100644 --- a/docs/src/models/tabM.md +++ b/docs/src/models/tabM.md @@ -1,5 +1,5 @@ # TabM -```@docs -TabMConfig +```@autodocs +Modules = [NeuroTabModels.Models.TabM] ``` From f92128707aab84641586d127665fbdc72b9da865 Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 9 Jun 2026 23:19:35 -0400 Subject: [PATCH 115/120] up --- src/models/MLP/mlp.jl | 26 ++++- src/models/MOETree/moetrees.jl | 50 +++++--- src/models/NeuroTree/model.jl | 17 +++ src/models/NeuroTree/neurotrees.jl | 50 +++++--- src/models/ResNet/resnet.jl | 28 ++++- src/models/TabM/TabM.jl | 39 ++++--- src/models/TabM/layers.jl | 134 ++++++++-------------- src/models/embeddings/batchnorm.jl | 24 +--- src/models/embeddings/compute_bins.jl | 16 +-- src/models/embeddings/config.jl | 35 +++--- src/models/embeddings/embeddings.jl | 9 +- src/models/embeddings/linear.jl | 21 ++-- src/models/embeddings/nlinear.jl | 23 ++-- src/models/embeddings/periodic.jl | 43 +++---- src/models/embeddings/piecewise_linear.jl | 43 +++---- src/models/models.jl | 7 ++ 16 files changed, 299 insertions(+), 266 deletions(-) diff --git a/src/models/MLP/mlp.jl b/src/models/MLP/mlp.jl index 380f12f..4e3c6e7 100644 --- a/src/models/MLP/mlp.jl +++ b/src/models/MLP/mlp.jl @@ -7,6 +7,18 @@ using LuxCore import ..Models: Architecture, get_activation +""" + MLPConfig(; kwargs...) + +Configuration for a multi-layer perceptron backbone. + +# Arguments +- `act::Symbol`: Activation — `:relu`, `:gelu`, `:sigmoid`, or `:tanh` (default `:relu`). +- `hidden_size::Int`: Hidden dimension (default `64`). +- `stack_size::Int`: Number of hidden blocks (default `1`). +- `dropout::Float64`: Dropout rate between blocks (default `0.0`). +- `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). +""" struct MLPConfig <: Architecture act::Symbol hidden_size::Int @@ -30,7 +42,7 @@ function MLPConfig(; kwargs...) args_default = setdiff(keys(args), keys(kwargs)) length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." + @info "Following $(length(args_default)) arguments set to default: $(join(args_default, ", "))." for arg in intersect(keys(args), keys(kwargs)) args[arg] = kwargs[arg] @@ -58,6 +70,18 @@ function _mlp_trunk(nfeats::Int, hsize::Int, outsize::Int, act, stack_size::Int, return Chain(layers...) end +""" + (config::MLPConfig)(; nfeats, outsize) + +Build a `Lux.Chain` from `config`. + +# Arguments +- `nfeats::Int`: Number of input features. +- `outsize::Int`: Number of output units. + +# Returns +A `Lux.Chain` of `Dense` → `BatchNorm` → `Dense` blocks with optional dropout. +""" function (config::MLPConfig)(; nfeats, outsize, kwargs...) act = get_activation(config.act) hsize = config.hidden_size diff --git a/src/models/MOETree/moetrees.jl b/src/models/MOETree/moetrees.jl index b1d243d..111968a 100644 --- a/src/models/MOETree/moetrees.jl +++ b/src/models/MOETree/moetrees.jl @@ -2,13 +2,12 @@ module MOETrees export MOETreeConfig -using Random using Lux using LuxCore +using Random: AbstractRNG using Statistics: mean -using NNlib: softplus, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh, tanhshrink +using NNlib: tanh_fast, hardtanh, tanhshrink -import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture include("model.jl") @@ -33,6 +32,23 @@ function StackedMOETree((ins, outs)::Pair{<:Integer,<:Integer}; hidden_size::Int return StackedMOETree(Chain(layers...)) end +""" + MOETreeConfig(; kwargs...) + +Configuration for mixture-of-experts tree ensembles. + +# Arguments +- `tree_type::Symbol`: `:binary` or `:oblivious` (default `:binary`). +- `actA::Symbol`: Feature activation — `:identity`, `:tanh`, `:hardtanh`, or `:tanhshrink` (default `:identity`). +- `depth::Int`: Tree depth (default `4`). +- `ntrees::Int`: Number of trees per layer (default `32`). +- `k::Int`: Output width multiplier (default `1`). +- `hidden_size::Int`: Hidden dimension for stacked trees (default `1`). +- `stack_size::Int`: Number of stacked tree layers (default `1`). +- `scaler::Bool`: Apply softplus scaling on tree logits (default `true`). +- `init_scale::Float32`: Leaf weight init scale (default `0.1`). +- `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). +""" struct MOETreeConfig <: Architecture tree_type::Symbol actA::Symbol @@ -66,10 +82,9 @@ function MOETreeConfig(; kwargs...) args_default = setdiff(keys(args), keys(kwargs)) length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." + @info "Following $(length(args_default)) arguments set to default: $(join(args_default, ", "))." - args_override = intersect(keys(args), keys(kwargs)) - for arg in args_override + for arg in intersect(keys(args), keys(kwargs)) args[arg] = kwargs[arg] end @@ -99,6 +114,18 @@ function _tree_kwargs(config::MOETreeConfig) ) end +""" + (config::MOETreeConfig)(; nfeats, outsize) + +Build a `Lux.Chain` from `config`. + +# Arguments +- `nfeats::Int`: Number of input features. +- `outsize::Int`: Number of output units. + +# Returns +A `Lux.Chain` of stacked mixture-of-experts tree layers. +""" function (config::MOETreeConfig)(; nfeats, outsize, kwargs...) kwargs = _tree_kwargs(config) @@ -138,13 +165,10 @@ function _tanhshrink_act(x) end """ - act_dict = Dict( - :identity => _identity_act, - :tanh => _tanh_act, - :hardtanh => _hardtanh_act, - :tanhshrink => _tanhshrink_act, - ) -Dictionary mapping features activation name to their function. + act_dict + +Dictionary mapping feature activation symbols to their functions. +Supported keys: `:identity`, `:tanh`, `:hardtanh`, `:tanhshrink`. """ const act_dict = Dict( :identity => _identity_act, diff --git a/src/models/NeuroTree/model.jl b/src/models/NeuroTree/model.jl index 8b476b2..9431eb0 100644 --- a/src/models/NeuroTree/model.jl +++ b/src/models/NeuroTree/model.jl @@ -1,3 +1,20 @@ +""" + NeuroTree(feats => outs; tree_type=:binary, actA=identity, scaler=true, + depth, trees, k, init_scale=0.1) + +Differentiable oblivious or binary tree ensemble layer. + +# Arguments +- `feats::Int`: Number of input features. +- `outs::Int`: Number of output targets per tree ensemble. +- `tree_type::Symbol`: `:binary` or `:oblivious`. +- `actA`: Feature activation applied to split weights. +- `scaler::Bool`: Scale logits with a learned softplus factor. +- `depth::Int`: Tree depth. +- `trees::Int`: Number of trees in the ensemble. +- `k::Int`: Output width multiplier averaged over trees. +- `init_scale::Float32`: Standard deviation for leaf weight initialization. +""" struct NeuroTree{F} <: AbstractLuxLayer tree_type::Symbol actA::F diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index 4c262ca..601ed7f 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -2,13 +2,12 @@ module NeuroTrees export NeuroTreeConfig -using Random using Lux using LuxCore +using Random: AbstractRNG using Statistics: mean -using NNlib: softplus, sigmoid_fast, hardsigmoid, tanh_fast, hardtanh, tanhshrink +using NNlib: tanh_fast, hardtanh, tanhshrink -import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture include("model.jl") @@ -33,6 +32,23 @@ function StackedNeuroTree((ins, outs)::Pair{<:Integer,<:Integer}; hidden_size::I return StackedNeuroTree(Chain(layers...)) end +""" + NeuroTreeConfig(; kwargs...) + +Configuration for differentiable neuro-tree ensembles. + +# Arguments +- `tree_type::Symbol`: `:binary` or `:oblivious` (default `:binary`). +- `actA::Symbol`: Feature activation — `:identity`, `:tanh`, `:hardtanh`, or `:tanhshrink` (default `:identity`). +- `depth::Int`: Tree depth (default `4`). +- `ntrees::Int`: Number of trees per layer (default `32`). +- `k::Int`: Output width multiplier (default `1`). +- `hidden_size::Int`: Hidden dimension for stacked trees (default `1`). +- `stack_size::Int`: Number of stacked tree layers (default `1`). +- `scaler::Bool`: Apply softplus scaling on tree logits (default `true`). +- `init_scale::Float32`: Leaf weight init scale (default `0.1`). +- `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). +""" struct NeuroTreeConfig <: Architecture tree_type::Symbol actA::Symbol @@ -66,10 +82,9 @@ function NeuroTreeConfig(; kwargs...) args_default = setdiff(keys(args), keys(kwargs)) length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." + @info "Following $(length(args_default)) arguments set to default: $(join(args_default, ", "))." - args_override = intersect(keys(args), keys(kwargs)) - for arg in args_override + for arg in intersect(keys(args), keys(kwargs)) args[arg] = kwargs[arg] end @@ -99,6 +114,18 @@ function _tree_kwargs(config::NeuroTreeConfig) ) end +""" + (config::NeuroTreeConfig)(; nfeats, outsize) + +Build a `Lux.Chain` from `config`. + +# Arguments +- `nfeats::Int`: Number of input features. +- `outsize::Int`: Number of output units. + +# Returns +A `Lux.Chain` of stacked neuro-tree layers. +""" function (config::NeuroTreeConfig)(; nfeats, outsize, kwargs...) kwargs = _tree_kwargs(config) @@ -138,13 +165,10 @@ function _tanhshrink_act(x) end """ - act_dict = Dict( - :identity => _identity_act, - :tanh => _tanh_act, - :hardtanh => _hardtanh_act, - :tanhshrink => _tanhshrink_act, - ) -Dictionary mapping features activation name to their function. + act_dict + +Dictionary mapping feature activation symbols to their functions. +Supported keys: `:identity`, `:tanh`, `:hardtanh`, `:tanhshrink`. """ const act_dict = Dict( :identity => _identity_act, diff --git a/src/models/ResNet/resnet.jl b/src/models/ResNet/resnet.jl index 66749ea..455db29 100644 --- a/src/models/ResNet/resnet.jl +++ b/src/models/ResNet/resnet.jl @@ -7,6 +7,18 @@ using LuxCore import ..Models: Architecture, get_activation +""" + ResNetConfig(; kwargs...) + +Configuration for a residual MLP backbone for tabular data. + +# Arguments +- `stack_size::Int`: Number of residual blocks (default `1`). +- `hidden_size::Int`: Hidden dimension (default `64`). +- `act::Symbol`: Activation — `:relu`, `:gelu`, `:sigmoid`, or `:tanh` (default `:relu`). +- `dropout::Float64`: Dropout rate within each block (default `0.0`). +- `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). +""" struct ResNetConfig <: Architecture stack_size::Int hidden_size::Int @@ -30,7 +42,7 @@ function ResNetConfig(; kwargs...) args_default = setdiff(keys(args), keys(kwargs)) length(args_default) > 0 && - @info "Following $(length(args_default)) arguments were not provided and will be set to default: $(join(args_default, ", "))." + @info "Following $(length(args_default)) arguments set to default: $(join(args_default, ", "))." for arg in intersect(keys(args), keys(kwargs)) args[arg] = kwargs[arg] @@ -72,6 +84,20 @@ function _resnet_trunk(nfeats::Int, hsize::Int, outsize::Int, act, stack_size::I return Chain(layers...) end +""" + (config::ResNetConfig)(; nfeats, outsize) + +Build a `Lux.Chain` from `config`. + +Each block applies two `Dense` layers with batch norm and a skip connection. + +# Arguments +- `nfeats::Int`: Number of input features. +- `outsize::Int`: Number of output units. + +# Returns +A `Lux.Chain` with residual blocks. +""" function (config::ResNetConfig)(; nfeats, outsize, kwargs...) act = get_activation(config.act) hsize = config.hidden_size diff --git a/src/models/TabM/TabM.jl b/src/models/TabM/TabM.jl index 2e1abc0..5932c0a 100644 --- a/src/models/TabM/TabM.jl +++ b/src/models/TabM/TabM.jl @@ -2,12 +2,11 @@ module TabM export TabMConfig -using Random using Lux using LuxCore -using NNlib +using LuxLib: batched_matmul +using Random: AbstractRNG, rand, randn -import ..Losses: get_loss_type, GaussianMLE import ..Models: Architecture, _broadcast_relu include("layers.jl") @@ -60,21 +59,15 @@ end """ TabMConfig(; kwargs...) -Configuration for TabM ensemble architectures. - -The chain output is 3D: `(outsize, k, batch)`. Ensemble averaging is handled -during training and at inference. +Configuration for TabM ensemble backbones. # Arguments -- `k::Int`: Number of ensemble members (default `32`). +- `k::Int`: Ensemble size (default `32`). - `n_blocks::Int`: Number of MLP blocks (default `3`). -- `d_block::Int`: Hidden dimension per block (default `512`). +- `d_block::Int`: Hidden dimension (default `512`). - `dropout::Float64`: Dropout rate (default `0.1`). - `arch_type::Symbol`: `:tabm`, `:tabm_mini`, or `:tabm_packed` (default `:tabm`). -- `scaling_init::Union{Nothing,Symbol}`: Init for ensemble scaling — `:random_signs`, `:normal`, - `:ones`, or `nothing` for auto-detection (default `nothing`). - When `nothing`, defaults to `:random_signs` without embeddings and `:normal` with embeddings. - Explicit values always take priority. +- `scaling_init::Union{Nothing,Symbol}`: `:random_signs`, `:normal`, `:ones`, or `nothing` (default `nothing`). - `MLE_tree_split::Bool`: Split output head for Gaussian MLE (default `false`). """ struct TabMConfig <: Architecture @@ -120,6 +113,17 @@ function TabMConfig(; kwargs...) ) end +""" + (config::TabMConfig)(; nfeats, outsize, d_features=nothing, scaling_init_override=nothing) + +Build a `Lux.Chain` from `config`. Output shape is `(outsize, k, batch)`. + +# Arguments +- `nfeats::Int`: Number of input features. +- `outsize::Int`: Number of output units. +- `d_features`: Per-feature sizes for grouped scaling init (default `ones(Int, nfeats)`). +- `scaling_init_override`: Used when `config.scaling_init` is `nothing`. +""" function (config::TabMConfig)(; nfeats, outsize, d_features=nothing, scaling_init_override=nothing) @assert config.k > 0 "k must be > 0, got $(config.k)" @assert nfeats > 0 "nfeats must be > 0, got $nfeats" @@ -159,11 +163,12 @@ function (config::TabMConfig)(; nfeats, outsize, d_features=nothing, scaling_ini error("Unknown arch_type: $(config.arch_type)") end - head = if config.MLE_tree_split && outsize == 2 - split_out = outsize ÷ 2 + head = if config.MLE_tree_split + iseven(outsize) || error("MLE_tree_split requires an even `outsize` (e.g., 2 for μ and σ). Got: $outsize") + head_outsize = outsize ÷ 2 Parallel(vcat, - LinearEnsemble(d_block, split_out, k), - LinearEnsemble(d_block, split_out, k)) + LinearEnsemble(d_block, head_outsize, k), + LinearEnsemble(d_block, head_outsize, k)) else LinearEnsemble(d_block, outsize, k) end diff --git a/src/models/TabM/layers.jl b/src/models/TabM/layers.jl index 695704a..f06d20c 100644 --- a/src/models/TabM/layers.jl +++ b/src/models/TabM/layers.jl @@ -1,66 +1,9 @@ -using Lux: AbstractLuxLayer, WrappedFunction -using LuxCore -using LuxLib: batched_matmul -using Random: AbstractRNG - -""" - _init_rsqrt_uniform(rng, dims, d) → Array{Float32} - -Uniform init in `[-1/√d, 1/√d]`. -""" -function _init_rsqrt_uniform(rng::AbstractRNG, dims, d::Int) - s = Float32(1 / sqrt(d)) - return s .* (2f0 .* rand(rng, Float32, dims...) .- 1f0) -end - -""" - _init_scaling(rng, dims, init) → Array{Float32} - -Initialize scaling vectors for ensemble adapters. -- `:ones` — deterministic ones -- `:normal` — N(0,1) -- `:random_signs` — ±1 -""" -function _init_scaling(rng::AbstractRNG, dims, init::Symbol) - if init == :ones - return ones(Float32, dims...) - elseif init == :normal - return randn(rng, Float32, dims...) - elseif init == :random_signs - return Float32.(2 .* (rand(rng, Float32, dims...) .> 0.5f0) .- 1) - else - error("Unknown scaling init: $init") - end -end - -""" - _init_scaling_with_chunks(rng, dims, init, chunks) → Array{Float32} - -Initialize scaling with grouped chunks. Each chunk shares the same random value -per ensemble member, providing structured diversity for features with different -representation sizes. -""" -function _init_scaling_with_chunks(rng::AbstractRNG, dims::Tuple{Int,Int}, - init::Symbol, chunks::Vector{Int}) - d, k = dims - @assert d == sum(chunks) "Chunks must sum to $d, got $(sum(chunks))" - weight = zeros(Float32, d, k) - row = 1 - for chunk_size in chunks - val = _init_scaling(rng, (1, k), init) - weight[row:row+chunk_size-1, :] .= repeat(val, chunk_size, 1) - row += chunk_size - end - return weight -end - """ EnsembleView(k) -Broadcasts a `(D, B)` matrix to `(D, K, B)` by repeating along dim 2. -Passes through `(D, K, B)` input unchanged. +Repeat `(D, B)` input to `(D, K, B)`. Passes through `(D, K, B)` unchanged. """ -struct EnsembleView <: AbstractLuxLayer +struct EnsembleView <: LuxCore.AbstractLuxLayer k::Int end @@ -78,21 +21,16 @@ end LinearBatchEnsemble(in_f, out_f; k, scaling_init=:random_signs, first_scaling_init_chunks=nothing, bias=true) -Batch-ensemble linear: `y = S ⊙ (W(R ⊙ x)) + bias` with shared `W` and -per-member `R`, `S`, `bias`. - -Equivalent to defining per-member weight matrices `Wᵢ = W ⊙ (sᵢrᵢᵀ)`. +Batch-ensemble linear: `y = S ⊙ (W(R ⊙ x)) + bias`. # Arguments -- `in_f::Int`: Input dimension. -- `out_f::Int`: Output dimension. -- `k::Int`: Number of ensemble members. -- `scaling_init`: Init for R/S. A `Symbol` (applied to both) or - `Tuple{Symbol,Symbol}` for `(R, S)`. Options: `:ones`, `:normal`, `:random_signs`. -- `first_scaling_init_chunks`: Chunk sizes for grouped R init. -- `bias::Bool`: Include per-member bias (default `true`). +- `in_f`, `out_f`: Input and output dimensions. +- `k::Int`: Ensemble size. +- `scaling_init`: `:ones`, `:normal`, or `:random_signs`; or `(R, S)` tuple. +- `first_scaling_init_chunks`: Grouped init for input scaling `R`. +- `bias::Bool`: Per-member bias (default `true`). """ -struct LinearBatchEnsemble <: AbstractLuxLayer +struct LinearBatchEnsemble <: LuxCore.AbstractLuxLayer in_features::Int out_features::Int k::Int @@ -143,15 +81,9 @@ end """ LinearEnsemble(in_f, out_f, k; bias=true) -`k` independent linear layers applied via `batched_matmul`. - -# Arguments -- `in_f::Int`: Input dimension. -- `out_f::Int`: Output dimension. -- `k::Int`: Number of ensemble members. -- `bias::Bool`: Include per-member bias (default `true`). +`k` independent linear layers via `batched_matmul`. Input/output `(features, k, batch)`. """ -struct LinearEnsemble <: AbstractLuxLayer +struct LinearEnsemble <: LuxCore.AbstractLuxLayer in_features::Int out_features::Int k::Int @@ -182,16 +114,9 @@ end """ ScaleEnsemble(k, d; init=:random_signs, init_chunks=nothing, bias=false) -Per-member elementwise scaling (and optional bias). - -# Arguments -- `k::Int`: Number of ensemble members. -- `d::Int`: Feature dimension. -- `init::Symbol`: Weight init — `:ones`, `:normal`, or `:random_signs`. -- `init_chunks`: Chunk sizes for grouped init. -- `bias::Bool`: Include per-member additive bias (default `false`). +Per-member elementwise scaling on `(d, k, batch)` input. """ -struct ScaleEnsemble <: AbstractLuxLayer +struct ScaleEnsemble <: LuxCore.AbstractLuxLayer k::Int d::Int init::Symbol @@ -226,4 +151,35 @@ function (m::ScaleEnsemble)(x::AbstractArray{T,3}, ps, st) where {T} else return x .* w, st end -end \ No newline at end of file +end + +function _init_rsqrt_uniform(rng::AbstractRNG, dims, d::Int) + s = Float32(1 / sqrt(d)) + return s .* (2f0 .* rand(rng, Float32, dims...) .- 1f0) +end + +function _init_scaling(rng::AbstractRNG, dims, init::Symbol) + if init == :ones + return ones(Float32, dims...) + elseif init == :normal + return randn(rng, Float32, dims...) + elseif init == :random_signs + return Float32.(2 .* (rand(rng, Float32, dims...) .> 0.5f0) .- 1) + else + error("Unknown scaling init: $init") + end +end + +function _init_scaling_with_chunks(rng::AbstractRNG, dims::Tuple{Int,Int}, + init::Symbol, chunks::Vector{Int}) + d, k = dims + @assert d == sum(chunks) "Chunks must sum to $d, got $(sum(chunks))" + weight = zeros(Float32, d, k) + row = 1 + for chunk_size in chunks + val = _init_scaling(rng, (1, k), init) + weight[row:row+chunk_size-1, :] .= repeat(val, chunk_size, 1) + row += chunk_size + end + return weight +end diff --git a/src/models/embeddings/batchnorm.jl b/src/models/embeddings/batchnorm.jl index fddf521..464215c 100644 --- a/src/models/embeddings/batchnorm.jl +++ b/src/models/embeddings/batchnorm.jl @@ -1,26 +1,14 @@ -using Lux -using Random -using NNlib - """ - BatchNormEmbeddings(n_features) - -Embeds each continuous feature via a learned affine transformation followed by -an activation: `activation(w_j * x_j + b_j)`. -Produces a `(d_embedding, n_features, batch)` tensor. + BatchNormEmbeddings(nfeats) -# Arguments -- `n_features::Int`: Number of input features. +Feature-wise `BatchNorm` on `(nfeats, batch)` input. Output shape matches input. """ -struct BatchNormEmbeddings{L} <: Lux.AbstractLuxWrapperLayer{:layer} +struct BatchNormEmbeddings{L} <: LuxCore.AbstractLuxWrapperLayer{:layer} layer::L end -function BatchNormEmbeddings(n_features::Int) - return BatchNormEmbeddings(BatchNorm(n_features)) -end +BatchNormEmbeddings(nfeats::Int) = BatchNormEmbeddings(BatchNorm(nfeats)) function (l::BatchNormEmbeddings)(x::AbstractMatrix, ps, st) - x_bn, st = l.layer(x, ps, st) - return x_bn, st -end \ No newline at end of file + return l.layer(x, ps, st) +end diff --git a/src/models/embeddings/compute_bins.jl b/src/models/embeddings/compute_bins.jl index d0afdfd..157d030 100644 --- a/src/models/embeddings/compute_bins.jl +++ b/src/models/embeddings/compute_bins.jl @@ -1,20 +1,16 @@ """ compute_bins(X; bins=48) -Compute quantile-based bin boundaries for `PiecewiseLinearEncoding`/`PiecewiseLinearEmbeddings`. +Quantile-based bin edges for piecewise-linear embeddings. -Note: `X` should have shape `(n_samples, n_features)`, which is the transpose of the -model input convention `(n_features, batch)`. Transpose your data before calling this. +`X` must have shape `(n_samples, n_features)` — the transpose of model input `(n_features, batch)`. # Arguments -- `X::AbstractMatrix`: Training data of shape `(n_samples, n_features)`. -- `bins::Union{Int, Vector{Int}}`: Number of bins per feature (default `48`). - A single `Int` applies the same count to all features. A `Vector{Int}` of length - `n_features` specifies per-feature bin counts. +- `X::AbstractMatrix`: Training data `(n_samples, n_features)`. +- `bins::Union{Int, Vector{Int}}`: Bin count per feature (default `48`). # Returns -- `Vector{Vector{Float32}}`: Bin edges for each feature. Each vector has between 2 and - `bins + 1` elements (fewer if quantiles coincide for low-cardinality features). +`Vector{Vector{Float32}}` of bin edges per feature. """ function compute_bins(X::AbstractMatrix; bins::Union{Int,Vector{Int}}=48) n_samples, n_features = size(X) @@ -45,4 +41,4 @@ function compute_bins(X::AbstractMatrix; bins::Union{Int,Vector{Int}}=48) bins[j] = edges end return bins -end \ No newline at end of file +end diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl index 8d902bc..cc78ef6 100644 --- a/src/models/embeddings/config.jl +++ b/src/models/embeddings/config.jl @@ -1,5 +1,3 @@ -using NNlib: relu, tanh, softplus, hardtanh, tanhshrink - const act_dict = Dict( :identity => identity, :relu => relu, @@ -12,22 +10,16 @@ const act_dict = Dict( """ EmbeddingConfig(; kwargs...) -Architecture-agnostic configuration for numerical feature embeddings. -Constructed by the user and passed as `embedding_config` to `NeuroTabRegressor`/`NeuroTabClassifier`. +Configuration for numerical feature embeddings. +Pass as `embedding_config` to `NeuroTabRegressor` / `NeuroTabClassifier`. # Arguments -- `embedding_type::Symbol`: `:periodic`, `:linear`, or `:piecewise` (default `:periodic`). -- `d_embedding::Int`: Embedding dimension per feature (default `24`). -- `activation::Symbol`: activation function symbol (default `:relu` for periodic/linear, - `:identity` for piecewise). -- `bins::Union{Int, Vector{Int}}`: Number of bins for piecewise embeddings (default `48`). -- `frequencies::Int`: Frequency components for periodic embeddings (default `48`). -- `frequencies_init_scale::Float32`: σ for periodic frequencies init (default `0.01f0`). - -# Callable - config(; nfeats, X_train=nothing)-Lux.Chain - -Returns a `Chain(embedding_layer, FlattenLayer())` ready to prepend to any backbone. +- `embedding_type::Symbol`: `:periodic`, `:linear`, `:piecewise`, or `:batchnorm` (default `:periodic`). +- `d_embedding::Int`: Embedding dimension per feature (default `16`; forced to `1` for `:batchnorm`). +- `activation::Symbol`: `:identity`, `:relu`, `:tanh`, etc. (default `:relu`, or `:identity` for `:piecewise`). +- `bins::Union{Int, Vector{Int}}`: Bin count for `:piecewise` (default `32`). +- `frequencies::Int`: Sinusoidal components for `:periodic` (default `32`). +- `frequencies_init_scale::Float32`: σ for periodic frequency init (default `0.01f0`). """ struct EmbeddingConfig embedding_type::Symbol @@ -48,13 +40,11 @@ function EmbeddingConfig(; ) embedding_type = Symbol(embedding_type) - # Default activation depends on embedding type if isnothing(activation) activation = embedding_type == :piecewise ? :identity : :relu end activation = Symbol(activation) - # Override d_embedding for batchnorm to 1 (for proper derivation of the number of input feature to core model block) if embedding_type == :batchnorm d_embedding = 1 end @@ -62,6 +52,15 @@ function EmbeddingConfig(; return EmbeddingConfig(embedding_type, d_embedding, activation, bins, frequencies, frequencies_init_scale) end +""" + (config::EmbeddingConfig)(; nfeats, x_train=nothing) + +Build `Chain(embedding, FlattenLayer())` from `config`. + +# Arguments +- `nfeats::Int`: Number of input features. +- `x_train`: Training matrix `(n_samples, n_features)`; required for `:piecewise`. +""" function (config::EmbeddingConfig)(; nfeats::Int, x_train=nothing) emb = if config.embedding_type == :periodic PeriodicEmbeddings(nfeats, config.d_embedding; diff --git a/src/models/embeddings/embeddings.jl b/src/models/embeddings/embeddings.jl index f97eb40..918f231 100644 --- a/src/models/embeddings/embeddings.jl +++ b/src/models/embeddings/embeddings.jl @@ -1,12 +1,11 @@ module Embeddings using Lux -using Lux: Chain, FlattenLayer +using Lux: BatchNorm, Chain, Dense, FlattenLayer using LuxCore -using Random using NNlib -using LuxLib: batched_matmul -import Statistics: quantile +using Random: AbstractRNG, rand, randn +using Statistics: quantile export NLinear, LinearEmbeddings export Periodic, PeriodicEmbeddings @@ -21,4 +20,4 @@ include("piecewise_linear.jl") include("batchnorm.jl") include("config.jl") -end \ No newline at end of file +end diff --git a/src/models/embeddings/linear.jl b/src/models/embeddings/linear.jl index 6c2bc65..da7d589 100644 --- a/src/models/embeddings/linear.jl +++ b/src/models/embeddings/linear.jl @@ -1,31 +1,24 @@ -using Lux -using Random -using NNlib - """ LinearEmbeddings(nfeats, d_embedding; activation=relu) -Embeds each continuous feature via a learned affine transformation followed by -an activation: `activation(w_j * x_j + b_j)`. -Produces a `(d_embedding, nfeats, batch)` tensor. +Per-feature affine map `activation(w * x + b)`. Output shape `(d_embedding, nfeats, batch)`. # Arguments - `nfeats::Int`: Number of input features. - `d_embedding::Int`: Embedding dimension per feature. -- `activation`: Activation function applied element-wise (default `relu`). - E.g. `relu`, `tanh`, `identity`. +- `activation`: Element-wise activation (default `relu`). """ -struct LinearEmbeddings{F} <: Lux.AbstractLuxLayer +struct LinearEmbeddings{F} <: LuxCore.AbstractLuxLayer nfeats::Int d_embedding::Int activation::F end -function LinearEmbeddings(nfeats::Int, d_embedding::Int; activation=NNlib.relu) +function LinearEmbeddings(nfeats::Int, d_embedding::Int; activation=relu) return LinearEmbeddings(nfeats, d_embedding, activation) end -function Lux.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) +function LuxCore.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) limit = Float32(l.d_embedding)^(-0.5f0) weight = reshape((rand(rng, Float32, l.d_embedding, l.nfeats) .* 2f0 .* limit) .- limit, l.d_embedding, l.nfeats, 1) @@ -34,9 +27,9 @@ function Lux.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) return (weight=weight, bias=bias) end -Lux.initialstates(::AbstractRNG, ::LinearEmbeddings) = (;) +LuxCore.initialstates(::AbstractRNG, ::LinearEmbeddings) = (;) function (l::LinearEmbeddings)(x::AbstractMatrix, ps, st) x_r = reshape(x, 1, size(x, 1), size(x, 2)) return l.activation.(muladd.(ps.weight, x_r, ps.bias)), st -end \ No newline at end of file +end diff --git a/src/models/embeddings/nlinear.jl b/src/models/embeddings/nlinear.jl index fc80d5f..d575edf 100644 --- a/src/models/embeddings/nlinear.jl +++ b/src/models/embeddings/nlinear.jl @@ -1,20 +1,15 @@ -using Lux -using Random -using NNlib - """ NLinear(n, in_features, out_features; bias=true) -A batch of `n` independent linear layers applied in parallel via `batched_mul`. -Input shape `(in_features, n, batch)` → output shape `(out_features, n, batch)`. +`n` independent linear layers on `(in_features, n, batch)` via `batched_mul`. # Arguments -- `n::Int`: Number of independent linear layers (typically one per feature). -- `in_features::Int`: Input dimension for each linear layer. -- `out_features::Int`: Output dimension for each linear layer. -- `bias::Bool`: Whether to include a learnable bias (default `true`). +- `n::Int`: Number of parallel linear layers. +- `in_features::Int`: Input dimension per layer. +- `out_features::Int`: Output dimension per layer. +- `bias::Bool`: Include bias (default `true`). """ -struct NLinear <: Lux.AbstractLuxLayer +struct NLinear <: LuxCore.AbstractLuxLayer n::Int in_features::Int out_features::Int @@ -25,7 +20,7 @@ function NLinear(n::Int, in_features::Int, out_features::Int; bias::Bool=true) return NLinear(n, in_features, out_features, bias) end -function Lux.initialparameters(rng::AbstractRNG, l::NLinear) +function LuxCore.initialparameters(rng::AbstractRNG, l::NLinear) limit = Float32(l.in_features)^(-0.5f0) weight = (rand(rng, Float32, l.out_features, l.in_features, l.n) .* 2f0 .* limit) .- limit @@ -36,7 +31,7 @@ function Lux.initialparameters(rng::AbstractRNG, l::NLinear) end end -Lux.initialstates(::AbstractRNG, ::NLinear) = (;) +LuxCore.initialstates(::AbstractRNG, ::NLinear) = (;) function (l::NLinear)(x::AbstractArray{T,3}, ps, st) where T x_perm = PermutedDimsArray(x, (1, 3, 2)) @@ -47,4 +42,4 @@ function (l::NLinear)(x::AbstractArray{T,3}, ps, st) where T end return permutedims(out, (1, 3, 2)), st -end \ No newline at end of file +end diff --git a/src/models/embeddings/periodic.jl b/src/models/embeddings/periodic.jl index 503f763..a96a144 100644 --- a/src/models/embeddings/periodic.jl +++ b/src/models/embeddings/periodic.jl @@ -1,32 +1,27 @@ -using Lux -using Random -using NNlib - """ Periodic(nfeats, n_frequencies, sigma) -Maps each feature to `2 * n_frequencies` sinusoidal components: `[cos(2π w x), sin(2π w x)]`. -Output shape `(2 * n_frequencies, nfeats, batch)`. +Sinusoidal encoding `[cos(2π w x), sin(2π w x)]`. Output shape `(2 * n_frequencies, nfeats, batch)`. # Arguments - `nfeats::Int`: Number of input features. -- `n_frequencies::Int`: Number of frequency components per feature. -- `sigma::Float32`: Std-dev for the frequency weight initialization (clamped to ±3σ). +- `n_frequencies::Int`: Frequency components per feature. +- `sigma::Float32`: Std-dev for frequency weight init (clamped to ±3σ). """ -struct Periodic <: Lux.AbstractLuxLayer +struct Periodic <: LuxCore.AbstractLuxLayer nfeats::Int n_frequencies::Int sigma::Float32 end -function Lux.initialparameters(rng::AbstractRNG, l::Periodic) +function LuxCore.initialparameters(rng::AbstractRNG, l::Periodic) bound = l.sigma * 3f0 w = clamp.(l.sigma .* randn(rng, Float32, l.n_frequencies, l.nfeats), -bound, bound) w = reshape(2f0 * Float32(π) .* w, l.n_frequencies, l.nfeats, 1) return (weight=w,) end -Lux.initialstates(::AbstractRNG, ::Periodic) = (;) +LuxCore.initialstates(::AbstractRNG, ::Periodic) = (;) function (l::Periodic)(x::AbstractMatrix, ps, st) x_r = reshape(x, 1, size(x, 1), size(x, 2)) @@ -35,22 +30,20 @@ function (l::Periodic)(x::AbstractMatrix, ps, st) end """ - PeriodicEmbeddings(nfeats, d_embedding=24; n_frequencies=48, + PeriodicEmbeddings(nfeats, d_embedding=24; frequencies=48, frequencies_init_scale=0.01f0, activation=relu, lite=false) -Periodic sinusoidal encoding followed by a learned linear projection. -Applies `Periodic` → `NLinear` (or `Dense` if `lite`) → activation. +`Periodic` followed by per-feature linear projection and activation. # Arguments - `nfeats::Int`: Number of input features. -- `d_embedding::Int`: Output embedding dimension per feature (default `24`). -- `n_frequencies::Int`: Sinusoidal frequency components per feature (default `48`). -- `frequencies_init_scale::Float32`: σ for frequency weight init (default `0.01f0`). -- `activation`: Activation function applied after projection (default `relu`). E.g. `relu`, `tanh`, `identity`. -- `lite::Bool`: Use a single shared `Dense` instead of per-feature `NLinear` (default `false`). - Only valid when `activation` is not `identity`. +- `d_embedding::Int`: Output dimension per feature (default `24`). +- `frequencies::Int`: Sinusoidal components per feature (default `48`). +- `frequencies_init_scale::Float32`: σ for frequency init (default `0.01f0`). +- `activation`: Post-projection activation (default `relu`). +- `lite::Bool`: Use shared `Dense` instead of per-feature `NLinear` (default `false`). """ -struct PeriodicEmbeddings{P,L,F} <: Lux.AbstractLuxContainerLayer{(:periodic, :linear)} +struct PeriodicEmbeddings{P,L,F} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :linear)} periodic::P linear::L activation::F @@ -62,7 +55,7 @@ function PeriodicEmbeddings( d_embedding::Int=24; frequencies::Int=48, frequencies_init_scale::Float32=0.01f0, - activation=NNlib.relu, + activation=relu, lite::Bool=false, ) if lite && activation === identity @@ -89,7 +82,5 @@ function (m::PeriodicEmbeddings)(x::AbstractMatrix, ps, st) m.linear(h, ps.linear, st.linear) end - h_lin = m.activation.(h_lin) - - return h_lin, (periodic=st_p, linear=st_l) -end \ No newline at end of file + return m.activation.(h_lin), (periodic=st_p, linear=st_l) +end diff --git a/src/models/embeddings/piecewise_linear.jl b/src/models/embeddings/piecewise_linear.jl index a051cc9..223a7a4 100644 --- a/src/models/embeddings/piecewise_linear.jl +++ b/src/models/embeddings/piecewise_linear.jl @@ -1,17 +1,13 @@ -using Lux -using Random -using NNlib - """ PiecewiseLinearEncoding(bins) -Non-trainable piecewise-linear encoding using precomputed bin edges. +Non-trainable piecewise-linear encoding from precomputed bin edges. Output shape `(max_n_bins, nfeats, batch)`. # Arguments - `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). """ -struct PiecewiseLinearEncoding <: Lux.AbstractLuxLayer +struct PiecewiseLinearEncoding <: LuxCore.AbstractLuxLayer bins::Vector{Vector{Float32}} nfeats::Int max_n_bins::Int @@ -26,9 +22,9 @@ function PiecewiseLinearEncoding(bins::Vector{<:AbstractVector}) return PiecewiseLinearEncoding(bins_f32, nfeats, max_n_bins) end -Lux.initialparameters(::AbstractRNG, ::PiecewiseLinearEncoding) = (;) +LuxCore.initialparameters(::AbstractRNG, ::PiecewiseLinearEncoding) = (;) -function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) +function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) M, N = l.max_n_bins, l.nfeats weight = zeros(Float32, M, N) @@ -40,9 +36,6 @@ function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) b = -bin_edges[1:end-1] ./ bin_width nb = length(bin_edges) - 1 - # Place the last bin's weight/bias at the end row; - # remaining bins fill rows 1:nb-1. Unused rows stay zero - # and are clamped to [0, 1] harmlessly. weight[end, i] = w[end] bias[end, i] = b[end] if nb > 1 @@ -51,7 +44,6 @@ function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) end end - # Pre-reshape to 3D to avoid per-call reshape in forward pass return ( weight=reshape(weight, M, N, 1), bias=reshape(bias, M, N, 1), @@ -59,7 +51,6 @@ function Lux.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) end function (l::PiecewiseLinearEncoding)(x::AbstractMatrix, ps, st) - # For each feature j and bin b: h[b,j,:] = clamp(w[b,j]*x[j,:] + bias[b,j], 0, 1) x_r = reshape(x, 1, size(x, 1), size(x, 2)) h = clamp.(muladd.(st.weight, x_r, st.bias), 0f0, 1f0) return h, st @@ -68,18 +59,17 @@ end """ PiecewiseLinearEmbeddings(bins, d_embedding; activation=identity, version=:B) -Learnable embeddings on top of `PiecewiseLinearEncoding`. -Version `:A`: PLE -> NLinear (with bias). -Version `:B`: PLE -> NLinear (zero-init, no bias) + LinearEmbeddings residual. -Output shape `(d_embedding, nfeats, batch)`. +Learnable projection on top of [`PiecewiseLinearEncoding`](@ref). +Version `:A`: encoding → `NLinear` (with bias). +Version `:B`: encoding → zero-init `NLinear` + [`LinearEmbeddings`](@ref) residual. # Arguments -- `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). -- `d_embedding::Int`: Embedding dimension per feature. -- `activation`: Activation function applied after projection (default `identity`). E.g. `relu`, `tanh`. +- `bins::Vector{<:AbstractVector}`: Bin edges from [`compute_bins`](@ref). +- `d_embedding::Int`: Output dimension per feature. +- `activation`: Post-projection activation (default `identity`). - `version::Symbol`: `:A` or `:B` (default `:B`). """ -struct PiecewiseLinearEmbeddings{L0,I,L,F} <: Lux.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)} +struct PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)} linear0::L0 encoding::I linear::L @@ -98,22 +88,21 @@ function PiecewiseLinearEmbeddings( max_n_bins = maximum(length(b) - 1 for b in bins) encoding = PiecewiseLinearEncoding(bins) - # Residual path uses raw affine output (no activation) linear0 = (version == :B) ? LinearEmbeddings(nfeats, d_embedding; activation=identity) : nothing linear = NLinear(nfeats, max_n_bins, d_embedding; bias=(version == :A)) return PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version) end -function Lux.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbeddings) - ps_l0 = m.linear0 === nothing ? nothing : Lux.initialparameters(rng, m.linear0) - ps_enc = Lux.initialparameters(rng, m.encoding) +function LuxCore.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbeddings) + ps_l0 = m.linear0 === nothing ? nothing : LuxCore.initialparameters(rng, m.linear0) + ps_enc = LuxCore.initialparameters(rng, m.encoding) if m.version == :B n = m.linear.n ps_lin = (weight=zeros(Float32, m.linear.out_features, m.linear.in_features, n),) else - ps_lin = Lux.initialparameters(rng, m.linear) + ps_lin = LuxCore.initialparameters(rng, m.linear) end return (linear0=ps_l0, encoding=ps_enc, linear=ps_lin) @@ -134,4 +123,4 @@ function (m::PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st) h_final = m.activation.(h_final) return h_final, (linear0=st_l0, encoding=st_enc, linear=st_lin) -end \ No newline at end of file +end diff --git a/src/models/models.jl b/src/models/models.jl index a1f3725..3e8e3d1 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -8,6 +8,13 @@ using ..Losses using Lux: Chain using NNlib +""" + Architecture + +Abstract supertype for backbone configuration objects. + +Subtypes are functors: call with `(config)(; nfeats, outsize)` to build a `Lux.Chain`. +""" abstract type Architecture end _broadcast_relu(x) = NNlib.relu.(x) From 179e40916e204abf45fd153b8cb3eefd011abb8b Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Tue, 9 Jun 2026 23:51:22 -0400 Subject: [PATCH 116/120] up --- docs/make.jl | 2 ++ docs/src/API.md | 3 ++- src/infer.jl | 37 +++++++++++++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 1 deletion(-) diff --git a/docs/make.jl b/docs/make.jl index 0b86a9f..56a44f7 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,9 +2,11 @@ using Documenter using DocumenterVitepress using NeuroTabModels +using DataFrames pages = [ "Quick start" => "quick-start.md", + "API" => "API.md", "Design" => "design.md", "Models" => [ "Interface" => "models/models.md", diff --git a/docs/src/API.md b/docs/src/API.md index dcbf626..0e38ca1 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -3,11 +3,12 @@ ## Training ```@docs -NeuroTabModels.fit +NeuroTabModels.Fit.fit ``` ## Inference ```@docs NeuroTabModels.infer +NeuroTabModels.NeuroTabModel(::AbstractDataFrame) ``` diff --git a/src/infer.jl b/src/infer.jl index 8bb0089..5a8370d 100644 --- a/src/infer.jl +++ b/src/infer.jl @@ -82,6 +82,22 @@ function _infer_grp_loop(::Val, chain, data, x0, dev, cdev, ps, st) return preds end +""" + infer(m::NeuroTabModel, data; device=:cpu, backend=get(m.info, :backend, :zygote), proj=true) + +Run inference on batched feature data. + +# Arguments + +- `m::NeuroTabModel`: A fitted model. +- `data`: Iterable of feature batches. + +# Keyword arguments + +- `device=:cpu`: Execution device (`:cpu` or `:gpu`). +- `backend`: AD backend (`:zygote`, `:enzyme`, or `:reactant`). Defaults to the value stored on the model. +- `proj=true`: When `true`, map raw outputs to natural scale. Set to `false` for raw model-scale predictions. +""" function infer(m::NeuroTabModel{L}, data; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) where {L} device = Symbol(device) backend = Symbol(backend) @@ -122,6 +138,22 @@ function infer_grp(m::NeuroTabModel{L}, data; device=:cpu, backend=get(m.info, : return _scaler(L, p, scalers) end +""" + infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj=true) + +Run inference on tabular data and return predictions. + +# Arguments + +- `m::NeuroTabModel`: A fitted model. +- `df`: Feature data as an `AbstractDataFrame`. + +# Keyword arguments + +- `device=:cpu`: Execution device (`:cpu` or `:gpu`). +- `backend`: AD backend (`:zygote`, `:enzyme`, or `:reactant`). Defaults to the value stored on the model. +- `proj=true`: When `true`, map raw outputs to natural scale. Set to `false` for raw model-scale outputs. +""" function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) group_key = m.info[:group_key] if isnothing(group_key) @@ -135,6 +167,11 @@ function infer(m::NeuroTabModel, df::AbstractDataFrame; device=:cpu, backend=get return p end +""" + (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj=true) + +Run inference on tabular data. +""" function (m::NeuroTabModel)(df::AbstractDataFrame; device=:cpu, backend=get(m.info, :backend, :zygote), proj::Bool=true) return infer(m, df; device, backend, proj) end From 14e7d38a49fa50268922b720ad59a8e0d875fa47 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Wed, 10 Jun 2026 18:40:00 -0400 Subject: [PATCH 117/120] sync main --- docs/make.jl | 1 - docs/src/API.md | 1 - src/models/NeuroTree/neurotrees.jl | 7 ------- 3 files changed, 9 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index 56a44f7..fd2806d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,7 +2,6 @@ using Documenter using DocumenterVitepress using NeuroTabModels -using DataFrames pages = [ "Quick start" => "quick-start.md", diff --git a/docs/src/API.md b/docs/src/API.md index 0e38ca1..2329844 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -10,5 +10,4 @@ NeuroTabModels.Fit.fit ```@docs NeuroTabModels.infer -NeuroTabModels.NeuroTabModel(::AbstractDataFrame) ``` diff --git a/src/models/NeuroTree/neurotrees.jl b/src/models/NeuroTree/neurotrees.jl index a890634..929fb56 100644 --- a/src/models/NeuroTree/neurotrees.jl +++ b/src/models/NeuroTree/neurotrees.jl @@ -39,17 +39,10 @@ Configuration for differentiable neuro-tree ensembles. # Arguments - `tree_type::Symbol`: `:binary` or `:oblivious` (default `:binary`). -<<<<<<< HEAD -- `actA::Symbol`: Feature activation — `:identity`, `:tanh`, `:hardtanh`, or `:tanhshrink` (default `:identity`). -- `depth::Int`: Tree depth (default `4`). -- `ntrees::Int`: Number of trees per layer (default `32`). -- `k::Int`: Output width multiplier (default `1`). -======= - `actA::Symbol`: Feature activation. One of `:identity`, `:tanh`, `:hardtanh`, or `:tanhshrink` (default `:identity`). - `depth::Int`: Tree depth (default `4`). - `ntrees::Int`: Number of trees per layer (default `32`). - `k::Int`: Ensemble size. ->>>>>>> 6a45bba51e56c414a7af341ffe79e79d54a4ae38 - `hidden_size::Int`: Hidden dimension for stacked trees (default `1`). - `stack_size::Int`: Number of stacked tree layers (default `1`). - `scaler::Bool`: Apply softplus scaling on tree logits (default `true`). From 4586f5f8f63c00be74bab9935ca287549373cb36 Mon Sep 17 00:00:00 2001 From: "jeremie.db" Date: Wed, 10 Jun 2026 18:58:16 -0400 Subject: [PATCH 118/120] sync main --- test/core.jl | 4 ++-- test/embedding.jl | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/test/core.jl b/test/core.jl index b8387e9..a0dc557 100644 --- a/test/core.jl +++ b/test/core.jl @@ -187,12 +187,12 @@ end dtrain = df[train_indices, :] deval = df[setdiff(1:nrow(df), train_indices), :] - arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=1, d_block=32, dropout=0.1, arch_type) + arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=1, d_block=64, dropout=0.1, arch_type) learner = NeuroTabClassifier(arch; embedding_config=Dict(:embedding_type => :batchnorm), nrounds=200, batchsize=32, - early_stopping_rounds=5, + early_stopping_rounds=10, lr=1e-2, ) diff --git a/test/embedding.jl b/test/embedding.jl index df6b98f..e98413c 100644 --- a/test/embedding.jl +++ b/test/embedding.jl @@ -62,22 +62,22 @@ end @testset "$embedding_type" for embedding_type in [:periodic, :linear, :piecewise] - embedding_config = Dict(:embedding_type => embedding_type, :d_embedding => 16) + embedding_config = Dict(:embedding_type => embedding_type, :d_embedding => 8) if embedding_type == :piecewise embedding_config[:bins] = 16 elseif embedding_type == :periodic embedding_config[:frequencies] = 16 end - arch = NeuroTabModels.TabMConfig(; k=1, n_blocks=1, d_block=32, dropout=0.0) + arch = NeuroTabModels.TabMConfig(; k=4, n_blocks=1, d_block=128, dropout=0.0) learner = NeuroTabClassifier(arch; nrounds=500, - batchsize=32, + batchsize=64, lr=1e-2, - early_stopping_rounds=5, + early_stopping_rounds=20, embedding_config) - m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names, print_every_n=5) + m = NeuroTabModels.fit(learner, dtrain; deval, target_name, feature_names, print_every_n=5); ptrain = [argmax(x) for x in eachrow(m(dtrain))] peval = [argmax(x) for x in eachrow(m(deval))] From f99d51aabb36fe6a22fe1d6a252e63960d7dbdac Mon Sep 17 00:00:00 2001 From: "jeremie.desgagne.bouchard" Date: Sat, 13 Jun 2026 17:36:40 -0400 Subject: [PATCH 119/120] up --- benchmarks/titanic-logloss.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/benchmarks/titanic-logloss.jl b/benchmarks/titanic-logloss.jl index 53ece98..d224bc0 100644 --- a/benchmarks/titanic-logloss.jl +++ b/benchmarks/titanic-logloss.jl @@ -85,6 +85,12 @@ learner = NeuroTabRegressor( device=:cpu, backend=:reactant ) +learner.embedding_config +# main / lux-new +# NeuroTabModels.Models.Embeddings.EmbeddingConfig(:batchnorm, 1, :relu, 32, 32, 0.01f0) +# mnca/temporal +# NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.BatchNormEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.BatchNormEmbeddings(), nothing) +# NeuroTabModels.Models.Embeddings.EmbeddingLayer{NeuroTabModels.Models.Embeddings.PiecewiseLinearEmbeddings, Nothing}(NeuroTabModels.Models.Embeddings.PiecewiseLinearEmbeddings(8, 16, :identity, :B), nothing) @time m = NeuroTabModels.fit( learner, From b3e08aaff0064c2d337c0007fa5edff9f2ee96e3 Mon Sep 17 00:00:00 2001 From: Aditya Pandey <147182147+AdityaPandeyCN@users.noreply.github.com> Date: Fri, 3 Jul 2026 00:06:38 +0530 Subject: [PATCH 120/120] Add temporal embedding and ModernNCA model (#49) * temporal embedding Signed-off-by: AdityaPandeyCN * cleanup imports Signed-off-by: AdityaPandeyCN * docstrings Signed-off-by: AdityaPandeyCN * temporal only fix Signed-off-by: AdityaPandeyCN * docstrings Signed-off-by: AdityaPandeyCN * temporal only setting fix Signed-off-by: AdityaPandeyCN * refactor(embeddings): IdentityEmbedding default, Parallel temporal branch, kwarg constructors Signed-off-by: AdityaPandeyCN * add embedding docs and use testmode Signed-off-by: AdityaPandeyCN * fix piecewise binning for constant features Signed-off-by: AdityaPandeyCN --------- Signed-off-by: AdityaPandeyCN --- docs/make.jl | 1 + docs/src/embeddings.md | 38 +++ ext/NeuroTabModelsReactantExt.jl | 7 +- src/Fit/callback.jl | 5 +- src/Fit/fit.jl | 25 +- src/learners.jl | 32 +- src/models/ModernNCA/modernnca.jl | 365 ++++++++++++++++++++ src/models/embeddings/batchnorm.jl | 15 +- src/models/embeddings/compute_bins.jl | 10 +- src/models/embeddings/config.jl | 393 +++++++++++++++++++--- src/models/embeddings/embeddings.jl | 17 +- src/models/embeddings/linear.jl | 23 +- src/models/embeddings/periodic.jl | 44 ++- src/models/embeddings/piecewise_linear.jl | 39 ++- src/models/embeddings/temporal.jl | 77 +++++ src/models/models.jl | 68 +++- 16 files changed, 1013 insertions(+), 146 deletions(-) create mode 100644 src/models/ModernNCA/modernnca.jl create mode 100644 src/models/embeddings/temporal.jl diff --git a/docs/make.jl b/docs/make.jl index fd2806d..6ad07d9 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -7,6 +7,7 @@ pages = [ "Quick start" => "quick-start.md", "API" => "API.md", "Design" => "design.md", + "Embeddings" => "embeddings.md", "Models" => [ "Interface" => "models/models.md", "MLP" => "models/mlp.md", diff --git a/docs/src/embeddings.md b/docs/src/embeddings.md index 5d4d933..bfdcc3b 100644 --- a/docs/src/embeddings.md +++ b/docs/src/embeddings.md @@ -1,5 +1,43 @@ # Embeddings +Embeddings transform raw input columns before they reach a model backbone. Numerical +embeddings operate column-wise, while temporal embeddings reserve one column for +time features and concatenate that branch with the remaining numerical features. + +## Available Embeddings + +- `IdentityEmbedding`: leaves numerical features unchanged. +- `LinearEmbeddings`: projects each numerical feature to a learned vector. +- `PeriodicEmbeddings`: expands each feature with learned sinusoidal terms before projection. +- `PiecewiseLinearEmbeddings`: computes feature bins from the training data and embeds the resulting piecewise-linear encoding. +- `BatchNormEmbeddings`: batch-normalizes raw numerical features without expanding them. +- `TemporalEmbeddings`: embeds one time column with Fourier features and an optional trend term. + +## Configuration + +Use `EmbeddingLayer` to combine a numerical embedding with an optional temporal +branch. A dictionary form is also available for model configuration dictionaries. + +```julia +EmbeddingLayer(PeriodicEmbeddings(; d_embedding=24)) + +EmbeddingLayer(; + num=PiecewiseLinearEmbeddings(; bins=32), + temp=TemporalEmbeddings(; index=1), +) + +EmbeddingLayer(Dict( + :embedding_type => :periodic, + :d_embedding => 24, + :temporal => Dict(:index => 1), +)) +``` + +Piecewise-linear and temporal embeddings need training data when the embedding chain is +built. Use `needs_x_train` to check that requirement for a spec. + +## API + ```@autodocs Modules = [NeuroTabModels.Models.Embeddings] ``` diff --git a/ext/NeuroTabModelsReactantExt.jl b/ext/NeuroTabModelsReactantExt.jl index 791cc74..c4e5c90 100644 --- a/ext/NeuroTabModelsReactantExt.jl +++ b/ext/NeuroTabModelsReactantExt.jl @@ -16,12 +16,17 @@ function _get_device(::Val{:reactant}, ::Val{D}; gpuID::Integer=0) where {D} return reactant_device() end +_same_shape(a::AbstractArray, b::AbstractArray) = size(a) == size(b) +_same_shape(a::Tuple, b::Tuple) = + length(a) == length(b) && all(_same_shape(ai, bi) for (ai, bi) in zip(a, b)) +_same_shape(_, _) = false + function _infer_loop(::Val{:reactant}, chain, data, x0, dev, cdev, ps, st) compiled = @compile _forward_reduce(chain, dev(x0), ps, st) preds = Vector{AbstractArray}() for x in data - if size(x) == size(x0) + if _same_shape(x, x0) pred = compiled(chain, dev(x), ps, st) else pred = Reactant.@jit _forward_reduce(chain, dev(x), ps, st) diff --git a/src/Fit/callback.jl b/src/Fit/callback.jl index 04b7f7b..5fbaef6 100644 --- a/src/Fit/callback.jl +++ b/src/Fit/callback.jl @@ -7,6 +7,7 @@ using ..Learners: LearnerTypes using ...Infer: reduce_pred, _get_device using ..Data: get_df_loader_train using ..Metrics +using ...Models using Lux: Training, testmode @@ -26,7 +27,8 @@ end function CallBack( config::LearnerTypes, df::AbstractDataFrame, - cache; + cache, + m; feature_names, target_name, weight_name=nothing, @@ -43,6 +45,7 @@ function CallBack( deval = get_df_loader_train(dfg; feature_names, target_name, weight_name, offset_name, scalers, batchsize, shuffle=false) |> dev ps, st = ts.parameters, testmode(ts.states) + deval = Models.eval_dataloader(m.chain, m.info, deval, dev, ps, st) d0 = first(deval) eval_compiled = _build_eval_step(ts.model, feval, d0, ps, st; reactant=config.backend == :reactant) diff --git a/src/Fit/fit.jl b/src/Fit/fit.jl index cf9e41b..a3beec5 100644 --- a/src/Fit/fit.jl +++ b/src/Fit/fit.jl @@ -9,7 +9,7 @@ using ..Losses using ..Metrics using ..Infer: reduce_pred, _get_device -import Random: Xoshiro +import Random: Xoshiro, default_rng import Statistics: mean, std import MLJModelInterface: fit import Optimisers: OptimiserChain, WeightDecay, NAdam, Adam @@ -73,19 +73,12 @@ function init( # Build chain: optional embeddings + architecture backbone embed_config = config.embedding_config - if isnothing(embed_config) - chain = config.arch(; nfeats, outsize) - else - if embed_config.embedding_type == :piecewise - x_train = Matrix{Float32}(df[:, feature_names]) - else - x_train = nothing - end - embed_chain = embed_config(; nfeats, x_train) - d_in = nfeats * embed_config.d_embedding - d_features = fill(embed_config.d_embedding, nfeats) - chain = Chain(embed_chain, config.arch(; nfeats=d_in, outsize, d_features, scaling_init_override=:normal)) - end + x_train = Models.Embeddings.needs_x_train(embed_config) ? Matrix{Float32}(df[:, feature_names]) : nothing + embed_chain = Models.Embeddings.build_embedding_chain(embed_config, nfeats; x_train) + d_in = Models.Embeddings.embedding_width(embed_chain, randn(Float32, nfeats, 2), default_rng()) + chain = Models.build_chain(config.arch, embed_chain; + nfeats, outsize, d_in, embedding_config=embed_config, loss_type=L) + info = Dict( :nrounds => 0, @@ -105,6 +98,8 @@ function init( rng = Xoshiro(config.seed) ps, st = Lux.setup(rng, m.chain) |> dev + data = Models.train_dataloader(config.arch, m, data, df; + feature_names, target_name, loss_type=L, scalers, batchsize, dev, rng) opt = OptimiserChain(NAdam(config.lr), WeightDecay(config.wd)) ts = Training.TrainState(m.chain, ps, st, opt) @@ -164,7 +159,7 @@ function fit( logger = nothing if !isnothing(deval) - cb = CallBack(config, deval, cache; feature_names, target_name, weight_name, offset_name, group_key) + cb = CallBack(config, deval, cache, m; feature_names, target_name, weight_name, offset_name, group_key) logger = init_logger(config) cb(logger, 0, cache[:train_state]) (verbosity > 0) && @info "Init training" metric = logger[:metrics][end] diff --git a/src/learners.jl b/src/learners.jl index 811e00a..849ea18 100644 --- a/src/learners.jl +++ b/src/learners.jl @@ -5,14 +5,19 @@ import MLJModelInterface: fit, update, predict, schema import Random using ..Models -using ..Models: EmbeddingConfig +using ..Models: AbstractEmbedding, IdentityEmbedding, EmbeddingLayer export NeuroTabRegressor, NeuroTabClassifier, LearnerTypes +_to_embedding(::Nothing) = IdentityEmbedding() +_to_embedding(e::AbstractEmbedding) = e +_to_embedding(d::AbstractDict) = EmbeddingLayer(d) +_to_embedding(x) = error("`embedding_config` must be `nothing`, an `AbstractDict`, or an `AbstractEmbedding`; got $(typeof(x)).") + mutable struct NeuroTabRegressor <: MMI.Deterministic loss::Symbol metric::Symbol arch::Architecture - embedding_config::Union{Nothing,EmbeddingConfig} + embedding_config::AbstractEmbedding nrounds::Int early_stopping_rounds::Int lr::Float32 @@ -47,8 +52,10 @@ A model type for constructing a NeuroTabRegressor, based on [NeuroTabModels.jl]( - `backend=:zygote`: Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`. - `device=:gpu`: Execution device. One of `:cpu` or `:gpu`. - `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. -- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. - E.g. `embedding_config=Dict(:embedding_type => :periodic, :d_embedding => 24)`. +- `embedding_config=nothing`: Optional numerical/temporal embeddings. Accepts `nothing` (no-op), + an `AbstractEmbedding` (e.g. `EmbeddingLayer(num=PeriodicEmbeddings(d_embedding=24))`), or an + `AbstractDict` selecting the type via `:embedding_type` (e.g. + `Dict(:embedding_type => :periodic, :d_embedding => 24)`). `backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD. `backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. @@ -194,11 +201,7 @@ function NeuroTabRegressor(arch::Architecture; kwargs...) @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU." end - # Build EmbeddingConfig from Dict if needed - embed = args[:embedding_config] - if embed isa AbstractDict - embed = EmbeddingConfig(; embed...) - end + embed = _to_embedding(args[:embedding_config]) config = NeuroTabRegressor( loss, @@ -231,7 +234,7 @@ mutable struct NeuroTabClassifier <: MMI.Probabilistic loss::Symbol metric::Symbol arch::Architecture - embedding_config::Union{Nothing,EmbeddingConfig} + embedding_config::AbstractEmbedding nrounds::Int early_stopping_rounds::Int lr::Float32 @@ -259,7 +262,8 @@ A model type for constructing a NeuroTabClassifier, based on [NeuroTabModels.jl] - `backend=:zygote`: Backend used by Lux. One of `:enzyme`, `:zygote`, or `:reactant`. - `device=:gpu`: Execution device. One of `:cpu` or `:gpu`. - `gpuID=0`: GPU device to use, only relevant if `device = :gpu`. `0` auto-selects. -- `embedding_config=nothing`: Optional `Dict` or `EmbeddingConfig` for numerical feature embeddings. +- `embedding_config=nothing`: Optional numerical/temporal embeddings. Accepts `nothing` (no-op), + an `AbstractEmbedding`, or an `AbstractDict` selecting the type via `:embedding_type`. `backend=:zygote` works on `:cpu` and `:gpu`; `backend=:reactant` works on `:cpu` and `:gpu` and uses Enzyme for AD. `backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. @@ -390,11 +394,7 @@ function NeuroTabClassifier(arch::Architecture; kwargs...) @warn "`backend=:enzyme` with `device=:gpu` is currently known to fail for some NeuroTabModels models. Prefer `backend=:zygote` on GPU, `backend=:reactant` for Reactant, or `backend=:enzyme` on CPU." end - # Build EmbeddingConfig from Dict if needed - embed = args[:embedding_config] - if embed isa AbstractDict - embed = EmbeddingConfig(; embed...) - end + embed = _to_embedding(args[:embedding_config]) config = NeuroTabClassifier( :mlogloss, diff --git a/src/models/ModernNCA/modernnca.jl b/src/models/ModernNCA/modernnca.jl new file mode 100644 index 0000000..f60524c --- /dev/null +++ b/src/models/ModernNCA/modernnca.jl @@ -0,0 +1,365 @@ +module ModernNCA + +export ModernNCAConfig + +import ..Models +import ..Models: Architecture, NeuroTabModel +import ..Losses: LossType, MSE, MAE, LogLoss, MLogLoss, GaussianMLE + +using Lux +using LuxCore +using NNlib: relu, softmax +using Random: randperm, AbstractRNG +using DataFrames: AbstractDataFrame, select +using CategoricalArrays + +""" + ModernNCAConfig(; d_embedding=128, n_blocks=2, d_block=256, + dropout=0.1, temperature=1.0, sample_rate=0.8, + max_candidates=8192, eps=1f-8) + +Hyperparameters for ModernNCA. Pass as the `arch` argument to `NeuroTabRegressor` +or `NeuroTabClassifier`. + +# Arguments +- `d_embedding`: encoder output dimension. +- `n_blocks`, `d_block`: number and hidden width of post-encoder MLP blocks. +- `dropout`: dropout rate inside each block; skipped when `<= 0`. +- `temperature`: softmax temperature on negative pairwise distances. +- `sample_rate`: fraction of the batch complement sampled as candidates per step + (Stochastic Neighborhood Sampling). `1.0` uses the full complement. +- `max_candidates`: cap on sampled training candidates per batch. Set `<= 0` + to disable the cap. +- `eps`: numerical floor in `sqrt` and as a `temperature` lower bound. +""" +struct ModernNCAConfig <: Architecture + d_embedding::Int + n_blocks::Int + d_block::Int + dropout::Float32 + temperature::Float32 + sample_rate::Float32 + max_candidates::Int + eps::Float32 +end + +""" + ModernNCAModel + +Lux wrapper holding the backbone encoder, config, output size, and loss type. +""" +struct ModernNCAModel{B,LT} <: LuxCore.AbstractLuxWrapperLayer{:backbone} + backbone::B + cfg::ModernNCAConfig + outsize::Int + loss_type::Type{LT} +end + +""" + ModernNCALoader + +Training iterator yielding `((x, cand_x, cand_y, y), y)` batches. The corpus is +moved to device once at construction; query rows follow a shuffled epoch +permutation and candidates are resampled each step. + +# Fields +- `full_x`, `full_y`: device corpus. +- `batchsize`: query rows per step. +- `n_cand`: candidate rows sampled per step (`0` when `batchsize == N`). +- `rng`: RNG for sampling. +- `dev`: device-transfer callable. +""" +struct ModernNCALoader{X,Y,R<:AbstractRNG,D} + full_x::X + full_y::Y + batchsize::Int + n_cand::Int + rng::R + dev::D +end + +function ModernNCAConfig(; + d_embedding::Int=128, n_blocks::Int=2, d_block::Int=256, + dropout::Real=0.1, temperature::Real=1.0, sample_rate::Real=0.8, + max_candidates::Int=8192, eps::Real=1f-8, +) + return ModernNCAConfig(d_embedding, n_blocks, d_block, + Float32(dropout), Float32(temperature), Float32(sample_rate), + max_candidates, Float32(eps)) +end + +""" + _backbone(cfg, d_in, embedding_layer) + +Build the ModernNCA encoder: embedding, linear, then n_blocks times (BN, Dense(relu), Dropout, Dense), then BN. +""" +function _backbone(cfg::ModernNCAConfig, d_in::Int, embedding_layer) + emb = isnothing(embedding_layer) ? WrappedFunction(identity) : embedding_layer + layers = Any[emb, Dense(d_in => cfg.d_embedding)] + for _ in 1:cfg.n_blocks + push!(layers, BatchNorm(cfg.d_embedding)) + push!(layers, Dense(cfg.d_embedding => cfg.d_block, relu)) + cfg.dropout > 0 && push!(layers, Dropout(cfg.dropout)) + push!(layers, Dense(cfg.d_block => cfg.d_embedding)) + end + cfg.n_blocks > 0 && push!(layers, BatchNorm(cfg.d_embedding)) + return Chain(layers...) +end + +function (cfg::ModernNCAConfig)(; nfeats, outsize, + loss_type::Type{<:LossType}=MSE, + embedding_layer=nothing) + return ModernNCAModel(_backbone(cfg, nfeats, embedding_layer), cfg, Int(outsize), loss_type) +end + +Base.length(l::ModernNCALoader) = fld(size(l.full_x, 2), l.batchsize) + +""" + _pairwise_dist(q, k, ϵ) -> Matrix + +`(num_keys, batch)` Euclidean distance matrix between `q` `(d, batch)` and `k` `(d, num_keys)`. +`ϵ` is added under `sqrt` for numerical stability. +""" +function _pairwise_dist(q::AbstractMatrix, k::AbstractMatrix, ϵ::Float32) + q2 = sum(abs2, q; dims=1) # (1, batch) + k2 = sum(abs2, k; dims=1) # (1, num_keys) + kq = k' * q # (num_keys, batch) + return sqrt.(max.(0f0, k2' .+ q2 .- 2f0 .* kq) .+ ϵ) +end + +_diag_inf(i, j, d) = ifelse(i == j, typemax(typeof(d)), d) + +""" + _mask_diag(d) -> Matrix + +Set the diagonal of distance matrix `d` to `typemax`, collapsing self-attention +weights to zero during training. +""" +_mask_diag(d::AbstractMatrix) = + _diag_inf.(reshape(1:size(d, 1), :, 1), reshape(1:size(d, 2), 1, :), d) + +_to_output(::Type{<:Union{MSE,MAE}}, α, cy, _) = reshape(α' * cy, 1, :) + +function _to_output(::Type{<:LogLoss}, α, cy, _) + p = clamp.(reshape(α' * cy, 1, :), 1f-6, 1f0 - 1f-6) + return log.(p ./ (1f0 .- p)) +end + +function _to_output(::Type{<:MLogLoss}, α, cy, outsize::Int) + oh = ((k, c) -> ifelse(k == c, 1f0, 0f0)).( + reshape(UInt32(1):UInt32(outsize), :, 1), reshape(cy, 1, :)) + return log.(clamp.(oh * α, 1f-7, Inf32)) +end + +""" + _nca_logits(m, zq, zk, cy; mask_self=false) -> Matrix + +Softmax attention over corpus keys `zk` for query embeddings `zq`. +When `mask_self=true`, diagonal entries are masked (training path). +""" +function _nca_logits(m::ModernNCAModel, zq, zk, cy; mask_self::Bool=false) + d = _pairwise_dist(zq, zk, m.cfg.eps) ./ max(m.cfg.temperature, m.cfg.eps) + mask_self && (d = _mask_diag(d)) + α = softmax(-d; dims=1) + return _to_output(m.loss_type, α, cy, m.outsize) +end + +""" + _encode_corpus(m, cx, ps, st; chunk=2048) -> Matrix + +Encode raw corpus `cx` `(d_in, N)` in chunks of `chunk` rows to bound peak +memory and stay within BatchNorm shape limits. Returns `zk` `(d_embedding, N)`. +""" +function _encode_corpus(m::ModernNCAModel, cx, ps, st; chunk::Int=2048) + n = size(cx, 2) + if n <= chunk + return m.backbone(cx, ps, st)[1] + end + z0, _ = m.backbone(cx[:, 1:chunk], ps, st) + zk = similar(z0, size(z0, 1), n) + zk[:, 1:chunk] .= z0 + s = chunk + 1 + while s <= n + e = min(s + chunk - 1, n) + zk[:, s:e] .= m.backbone(cx[:, s:e], ps, st)[1] + s = e + 1 + end + return zk +end + +""" + (m::ModernNCAModel)((x, cand_x, cand_y, y), ps, st) + +Training forward: encodes queries and candidates separately, matching the +reference ModernNCA BatchNorm behavior. The encoded queries are prepended to +the candidate embeddings before NCA attention, then self-neighbors are masked. +""" +function (m::ModernNCAModel)((x, cand_x, cand_y, y)::Tuple{Any,Any,Any,Any}, ps, st) + zq, st_bb = m.backbone(x, ps, st) + zc, st_bb = size(cand_x, 2) == 0 ? + (similar(zq, size(zq, 1), 0), st_bb) : + m.backbone(cand_x, ps, st_bb) + z_all = hcat(zq, zc) + cy = vcat(vec(y), vec(cand_y)) + return _nca_logits(m, zq, z_all, cy; mask_self=true), st_bb +end + +""" + (m::ModernNCAModel)((x, cx, cy), ps, st) + +Inference/eval forward: encodes queries and the raw corpus `cx` with the current +parameters, then attends over the corpus. Encoding is done here (not pre-computed) +so the corpus embeddings always match the current model during evaluation. +""" +function (m::ModernNCAModel)((x, cand_x, cand_y)::Tuple{Any,Any,Any}, ps, st) + zq, st_bb = m.backbone(x, ps, st) + zk = _encode_corpus(m, cand_x, ps, st_bb) + return _nca_logits(m, zq, zk, cand_y), st_bb +end + +""" + Base.iterate(l::ModernNCALoader, state=nothing) + +`state = (perm, start)` carries the epoch permutation across calls, so query +batches are random without depending on dataframe row order. + +Candidates are sampled from the complement of the query batch within `perm`. +The index-skip mapping avoids allocating that complement: for `j` in +`1:(n-batchsize)`, use `perm[j]` before the batch window and +`perm[j + batchsize]` after it. `n_cand == 0` only when `batchsize == n`; then +the forward pass keys on the batch itself with diagonal masking. +""" +function Base.iterate(l::ModernNCALoader, state=nothing) + n = size(l.full_x, 2) + perm, start = state === nothing ? (randperm(l.rng, n), 1) : state + stop = start + l.batchsize - 1 + stop > n && return nothing + + batch_idx = perm[start:stop] + + bidx = l.dev(batch_idx) + x = l.full_x[:, bidx] + y = l.full_y[bidx] + + if l.n_cand > 0 + m = n - l.batchsize + js = rand(l.rng, 1:m, l.n_cand) + cidx = l.dev(perm[@. ifelse(js < start, js, js + l.batchsize)]) + cand_x = l.full_x[:, cidx] + cand_y = l.full_y[cidx] + else + cand_x = similar(l.full_x, size(l.full_x, 1), 0) + cand_y = similar(l.full_y, 0) + end + return ((x, cand_x, cand_y, y), y), (perm, stop + 1) +end + +""" + build_corpus(df, feature_names, target_name, loss_type, scalers) + +Return `(full_x, full_y)`: the corpus feature matrix `(d_in, N)` as `Float32` +and the encoded target vector, ready for `ModernNCALoader`. +""" +function build_corpus(df::AbstractDataFrame, feature_names, target_name, + loss_type::Type{<:LossType}, scalers) + full_x = permutedims(Matrix{Float32}(select(df, collect(feature_names)))) + full_y = _encode_targets(df, target_name, loss_type, scalers) + return full_x, full_y +end + +""" + _encode_targets(df, target_name, loss_type, scalers) + +Encode targets for each loss type: +- `MLogLoss`: 1-based `UInt32` class codes. +- `LogLoss`: `Float32` in `{0, 1}`. +- `MSE / MAE`: `Float32`, standardised when `scalers` is provided. + +# Arguments +- `df`: source data frame. +- `target_name`: target column. +- `loss_type`: target encoding dispatch. +- `scalers`: optional target scaler `(mu, sigma)`. +""" +_encode_targets(df, target_name, ::Type{<:MLogLoss}, _) = + UInt32.(CategoricalArrays.levelcode.(df[!, target_name])) + +function _encode_targets(df, target_name, ::Type{<:LogLoss}, _) + col = df[!, target_name] + eltype(col) <: CategoricalValue || return Float32.(col) + levels = CategoricalArrays.levels(col) + length(levels) == 2 || error("For `loss=:logloss`, target must have exactly 2 classes.") + return Float32.(CategoricalArrays.levelcode.(col) .- 1) +end + +function _encode_targets(df, target_name, ::Type{<:Union{MSE,MAE,GaussianMLE}}, scalers) + y = Float32.(df[!, target_name]) + isnothing(scalers) || (y .= (y .- scalers.mu) ./ scalers.sigma) + return y +end + +""" + Models.build_chain(cfg::ModernNCAConfig, embed_chain; outsize, d_in, loss_type, kw...) + +Pass the embedding layer into the backbone instead of wrapping it in an outer `Chain`. +""" +Models.build_chain(cfg::ModernNCAConfig, embed_chain; + outsize, d_in, loss_type, kw...) = + cfg(; nfeats=d_in, outsize, loss_type, embedding_layer=embed_chain) + +""" + Models.train_dataloader(cfg::ModernNCAConfig, ...) + +Build `ModernNCALoader` and stash the raw corpus on `m.info[:nca_ref]`. +`n_cand = floor(sample_rate × (N − batchsize))`, capped by +`cfg.max_candidates` when positive. `sample_rate ≥ 1` uses the full complement +before the cap; `batchsize == N` gives `n_cand = 0`. + +# Arguments +- `cfg`: ModernNCA config. +- `m`: fitted model wrapper. +- `df`: training data frame. +- `feature_names`, `target_name`: data columns. +- `loss_type`, `scalers`: target encoding metadata. +- `batchsize`, `dev`, `rng`: loader settings. +""" +function Models.train_dataloader(cfg::ModernNCAConfig, m::NeuroTabModel, ::Any, df; + feature_names, target_name, loss_type, scalers, batchsize, dev, rng, kw...) + cx, cy = build_corpus(df, feature_names, target_name, loss_type, scalers) + m.info[:nca_ref] = (cx=cx, cy=cy) + n = size(cx, 2) + batchsize = min(batchsize, n) + pool = n - batchsize + raw_n_cand = pool == 0 ? 0 : + cfg.sample_rate >= 1f0 ? pool : + max(Int(floor(cfg.sample_rate * pool)), 1) + n_cand = cfg.max_candidates > 0 ? min(raw_n_cand, cfg.max_candidates) : raw_n_cand + return ModernNCALoader(dev(cx), dev(cy), batchsize, n_cand, rng, dev) +end + +""" + Models.infer_dataloader(m::ModernNCAModel, ...) + +Wrap each inference batch as `(x, cx, cy)` with the raw training corpus. +The corpus is encoded inside the forward pass using the current parameters. +""" +function Models.infer_dataloader(::ModernNCAModel, info, data, dev, ps=nothing, st=nothing) + ref = info[:nca_ref] + cx, cy = dev(ref.cx), dev(ref.cy) + return Iterators.map(x -> (x, cx, cy), data) +end + +""" + Models.eval_dataloader(m::ModernNCAModel, ...) + +Wrap each eval batch as `((x, cx, cy), rest...)` with the raw training corpus. +The corpus is encoded inside the forward pass so embeddings always use the +current model parameters rather than a stale pre-encoded version. +""" +function Models.eval_dataloader(::ModernNCAModel, info, data, dev, ps=nothing, st=nothing) + ref = info[:nca_ref] + cx, cy = dev(ref.cx), dev(ref.cy) + return Iterators.map(d -> ((d[1], cx, cy), Base.tail(d)...), data) +end + +end # module diff --git a/src/models/embeddings/batchnorm.jl b/src/models/embeddings/batchnorm.jl index 464215c..538a96a 100644 --- a/src/models/embeddings/batchnorm.jl +++ b/src/models/embeddings/batchnorm.jl @@ -1,14 +1,17 @@ """ - BatchNormEmbeddings(nfeats) + _BatchNormEmbeddings(; nfeats) Feature-wise `BatchNorm` on `(nfeats, batch)` input. Output shape matches input. + +# Arguments +- `nfeats::Int`: Number of input features. """ -struct BatchNormEmbeddings{L} <: LuxCore.AbstractLuxWrapperLayer{:layer} +struct _BatchNormEmbeddings{L} <: LuxCore.AbstractLuxWrapperLayer{:layer} layer::L end -BatchNormEmbeddings(nfeats::Int) = BatchNormEmbeddings(BatchNorm(nfeats)) +_BatchNormEmbeddings(; nfeats::Int) = _BatchNormEmbeddings(BatchNorm(nfeats)) -function (l::BatchNormEmbeddings)(x::AbstractMatrix, ps, st) - return l.layer(x, ps, st) -end +(l::_BatchNormEmbeddings)(x::AbstractMatrix, ps, st) = l.layer(x, ps, st) + +LuxCore.outputsize(l::_BatchNormEmbeddings, x, ::AbstractRNG) = (size(x, 1),) \ No newline at end of file diff --git a/src/models/embeddings/compute_bins.jl b/src/models/embeddings/compute_bins.jl index 157d030..48e69a5 100644 --- a/src/models/embeddings/compute_bins.jl +++ b/src/models/embeddings/compute_bins.jl @@ -3,7 +3,7 @@ Quantile-based bin edges for piecewise-linear embeddings. -`X` must have shape `(n_samples, n_features)` — the transpose of model input `(n_features, batch)`. +`X` must have shape `(n_samples, n_features)`,the transpose of model input `(n_features, batch)`. # Arguments - `X::AbstractMatrix`: Training data `(n_samples, n_features)`. @@ -37,7 +37,13 @@ function compute_bins(X::AbstractMatrix; bins::Union{Int,Vector{Int}}=48) quantile_probs = range(0f0, 1f0, length=n_bins_vec[j] + 1) edges = Float32[quantile(col_buf, p; sorted=true) for p in quantile_probs] unique!(edges) - @assert length(edges) >= 2 "Feature $j has fewer than 2 unique bin edges" + if length(edges) < 2 + # Constant features collapse all quantile edges to one value. + # Give the encoder a tiny nonzero-width bin instead of failing. + v = edges[1] + delta = max(abs(v) * 1f-3, 1f-3) + edges = Float32[v - delta, v + delta] + end bins[j] = edges end return bins diff --git a/src/models/embeddings/config.jl b/src/models/embeddings/config.jl index cc78ef6..bb38a8b 100644 --- a/src/models/embeddings/config.jl +++ b/src/models/embeddings/config.jl @@ -7,77 +7,372 @@ const act_dict = Dict( :tanhshrink => tanhshrink, ) +# Supertype for embedding specs the learner can hold (`EmbeddingLayer`, numerical embeddings). +abstract type AbstractEmbedding end + +# Column-wise numerical-feature embeddings; also the type of an `EmbeddingLayer`'s `num`. +abstract type AbstractNumericalEmbedding <: AbstractEmbedding end + +# Time-column embeddings. +abstract type AbstractTemporalEmbedding end + +""" + IdentityEmbedding() + +No-op numerical embedding that passes features through unchanged; the default `num` for +an [`EmbeddingLayer`](@ref). +""" +struct IdentityEmbedding <: AbstractNumericalEmbedding end + """ - EmbeddingConfig(; kwargs...) + LinearEmbeddings(; d_embedding=16, activation=:relu) -Configuration for numerical feature embeddings. -Pass as `embedding_config` to `NeuroTabRegressor` / `NeuroTabClassifier`. +Per-feature linear embedding: each feature is projected to `d_embedding` by its own affine map, followed by `activation`. # Arguments -- `embedding_type::Symbol`: `:periodic`, `:linear`, `:piecewise`, or `:batchnorm` (default `:periodic`). -- `d_embedding::Int`: Embedding dimension per feature (default `16`; forced to `1` for `:batchnorm`). -- `activation::Symbol`: `:identity`, `:relu`, `:tanh`, etc. (default `:relu`, or `:identity` for `:piecewise`). -- `bins::Union{Int, Vector{Int}}`: Bin count for `:piecewise` (default `32`). -- `frequencies::Int`: Sinusoidal components for `:periodic` (default `32`). -- `frequencies_init_scale::Float32`: σ for periodic frequency init (default `0.01f0`). -""" -struct EmbeddingConfig - embedding_type::Symbol +- `d_embedding::Int`: Output dimension per feature (default `16`). +- `activation::Symbol`: Activation applied after the projection (default `:relu`). +""" +struct LinearEmbeddings <: AbstractNumericalEmbedding d_embedding::Int activation::Symbol - bins::Union{Int,Vector{Int}} +end +function LinearEmbeddings(; d_embedding::Int=16, activation::Symbol=:relu) + d_embedding > 0 || throw(ArgumentError("d_embedding must be > 0, got $d_embedding")) + LinearEmbeddings(d_embedding, activation) +end + +""" + PeriodicEmbeddings(; d_embedding=16, frequencies=32, frequencies_init_scale=0.01f0, + activation=:relu, lite=false) + +Per-feature periodic embedding: each feature is expanded with learned sine/cosine components, then projected to `d_embedding` and passed through `activation`. + +# Arguments +- `d_embedding::Int`: Output dimension per feature (default `16`). +- `frequencies::Int`: Number of sinusoidal components per feature (default `32`). +- `frequencies_init_scale::Float32`: Std. dev. initializing the frequencies (default `0.01f0`). +- `activation::Symbol`: Activation applied after the projection (default `:relu`). +- `lite::Bool`: Share the projection across features to cut parameters (default `false`). +""" +struct PeriodicEmbeddings <: AbstractNumericalEmbedding + d_embedding::Int frequencies::Int frequencies_init_scale::Float32 + activation::Symbol + lite::Bool end +function PeriodicEmbeddings(; d_embedding::Int=16, frequencies::Int=32, + frequencies_init_scale::Real=0.01f0, activation::Symbol=:relu, + lite::Bool=false) + d_embedding > 0 || throw(ArgumentError("d_embedding must be > 0, got $d_embedding")) + frequencies > 0 || throw(ArgumentError("frequencies must be > 0, got $frequencies")) + PeriodicEmbeddings(d_embedding, frequencies, Float32(frequencies_init_scale), + activation, lite) +end + +""" + PiecewiseLinearEmbeddings(; d_embedding=16, bins=32, activation=:identity, version=:B) -function EmbeddingConfig(; - embedding_type=:periodic, +Per-feature piecewise-linear embedding against bin edges computed from the training data, then projected to `d_embedding`. The bin edges are derived at fit time, so this embedding requires `x_train`. + +# Arguments +- `d_embedding::Int`: Output dimension per feature (default `16`). +- `bins::Union{Int,Vector{Int}}`: Number of bins, or per-feature bin counts (default `32`). +- `activation::Symbol`: Activation applied after the projection (default `:identity`). +- `version::Symbol`: Encoding variant, `:A` or `:B` (default `:B`). +""" +struct PiecewiseLinearEmbeddings <: AbstractNumericalEmbedding + d_embedding::Int + bins::Union{Int,Vector{Int}} + activation::Symbol + version::Symbol +end +function PiecewiseLinearEmbeddings(; d_embedding::Int=16, bins::Union{Int,Vector{Int}}=32, + activation::Symbol=:identity, version::Symbol=:B) + d_embedding > 0 || throw(ArgumentError("d_embedding must be > 0, got $d_embedding")) + version in (:A, :B) || + throw(ArgumentError("version must be :A or :B, got :$version")) + PiecewiseLinearEmbeddings(d_embedding, bins, activation, version) +end + +""" + BatchNormEmbeddings() + +Batch-normalize the raw features without expanding their dimension (one output per feature). +""" +struct BatchNormEmbeddings <: AbstractNumericalEmbedding end + +""" + TemporalEmbeddings(; index, order=[4, 1, 7, 0], periods=_DEFAULT_TEMPORAL_PERIODS, + trend=true, d_embedding=16) + +Fourier embedding of a single time column: the column at `index` is expanded into multi-scale sine/cosine features at `periods`, projected to `d_embedding`, and optionally augmented with a linear trend. `order` and `periods` align per band `order[i]` is the harmonic count for `periods[i]`. + +# Arguments +- `index::Int`: Required. 1-based position of the time column in `feature_names`. +- `order::Vector{Int}`: Harmonics per period; nonnegative with at least one positive entry (default `[4, 1, 7, 0]`). +- `periods::Vector{Float32}`: Base periods in column units (default `_DEFAULT_TEMPORAL_PERIODS`). +- `trend::Bool`: Append a linear trend term to the embedding (default `true`). +- `d_embedding::Int`: Projection dimension of the periodic features (default `16`). +""" +struct TemporalEmbeddings <: AbstractTemporalEmbedding + index::Int + order::Vector{Int} + periods::Vector{Float32} + trend::Bool + d_embedding::Int +end +function TemporalEmbeddings(; + index::Int, + order::AbstractVector{<:Integer}=Int[4, 1, 7, 0], + periods::AbstractVector{<:Real}=_DEFAULT_TEMPORAL_PERIODS, + trend::Bool=true, d_embedding::Int=16, - activation=nothing, - bins::Union{Int,Vector{Int}}=32, - frequencies::Int=32, - frequencies_init_scale::Float32=0.01f0, +) + index >= 1 || + throw(ArgumentError("index must be >= 1, got $index")) + length(order) == length(periods) || + throw(ArgumentError("length(order)=$(length(order)) must equal length(periods)=$(length(periods))")) + all(>=(0), order) && any(>(0), order) || + throw(ArgumentError("order must be non-negative with at least one positive entry")) + d_embedding > 0 || + throw(ArgumentError("d_embedding must be > 0, got $d_embedding")) + TemporalEmbeddings(index, Vector{Int}(order), Vector{Float32}(periods), + trend, d_embedding) +end + +""" + EmbeddingLayer(; num=IdentityEmbedding(), temp=nothing) + EmbeddingLayer(num::AbstractNumericalEmbedding; temp=nothing) + EmbeddingLayer(d::AbstractDict) + +Numerical embedding over the non-time columns, optionally combined with a temporal +embedding on the column at `temp.index`. `num` always exists (defaults to +[`IdentityEmbedding`](@ref), raw pass-through); `nothing` normalizes to it. When `temp` +is set the branches are concatenated features-first, temporal-last. + +The `Dict` form is for the hyperparameter harness: `:embedding_type` (`:linear`, +`:periodic`, `:piecewise`, `:batchnorm`, `:identity`) selects the numerical type, an +optional `:temporal` Dict gives [`TemporalEmbeddings`](@ref) kwargs, and unknown or +`nothing` keys are ignored (so a superset Dict works; missing `:embedding_type` → identity). + +```julia +EmbeddingLayer(PeriodicEmbeddings(; d_embedding=24)) +EmbeddingLayer(; temp=TemporalEmbeddings(; index=1)) # temporal only +EmbeddingLayer(Dict(:embedding_type => :periodic, :temporal => Dict(:index => 1))) +``` +""" +struct EmbeddingLayer{ + N<:AbstractNumericalEmbedding, + T<:Union{Nothing,AbstractTemporalEmbedding}, +} <: AbstractEmbedding + num::N + temp::T +end +_as_num(::Nothing) = IdentityEmbedding() +_as_num(n::AbstractNumericalEmbedding) = n +EmbeddingLayer(; num=IdentityEmbedding(), temp=nothing) = EmbeddingLayer(_as_num(num), temp) +EmbeddingLayer(num::AbstractNumericalEmbedding; temp=nothing) = EmbeddingLayer(num, temp) + +const _NUM_EMBEDDING_TYPES = Dict{Symbol,Type}( + :identity => IdentityEmbedding, + :linear => LinearEmbeddings, + :periodic => PeriodicEmbeddings, + :piecewise => PiecewiseLinearEmbeddings, + :batchnorm => BatchNormEmbeddings, ) - embedding_type = Symbol(embedding_type) - if isnothing(activation) - activation = embedding_type == :piecewise ? :identity : :relu +# Used by the `_num_from_dict` methods below. Filters a (possibly superset) config +# Dict down to the keyword arguments one embedding constructor accepts; absent and +# `nothing`-valued keys are skipped so the constructor's own defaults apply. +function _pick(d, keys) + ps = Pair{Symbol,Any}[] + for k in keys + haskey(d, k) && d[k] !== nothing && push!(ps, k => d[k]) end - activation = Symbol(activation) + return ps +end - if embedding_type == :batchnorm - d_embedding = 1 +_num_from_dict(::Type{IdentityEmbedding}, d) = IdentityEmbedding() +_num_from_dict(::Type{LinearEmbeddings}, d) = + LinearEmbeddings(; _pick(d, (:d_embedding, :activation))...) +_num_from_dict(::Type{PeriodicEmbeddings}, d) = + PeriodicEmbeddings(; _pick(d, (:d_embedding, :frequencies, :frequencies_init_scale, :activation, :lite))...) +_num_from_dict(::Type{PiecewiseLinearEmbeddings}, d) = + PiecewiseLinearEmbeddings(; _pick(d, (:d_embedding, :bins, :activation, :version))...) +_num_from_dict(::Type{BatchNormEmbeddings}, d) = BatchNormEmbeddings() + +function EmbeddingLayer(d::AbstractDict) + d = Dict{Symbol,Any}(Symbol(k) => v for (k, v) in d) + num = if haskey(d, :embedding_type) && d[:embedding_type] !== nothing + etype = Symbol(d[:embedding_type]) + T = get(_NUM_EMBEDDING_TYPES, etype) do + throw(ArgumentError("unknown :embedding_type $(repr(etype)); valid: $(collect(keys(_NUM_EMBEDDING_TYPES)))")) + end + _num_from_dict(T, d) + else + IdentityEmbedding() end + temp = if haskey(d, :temporal) + td = Dict{Symbol,Any}(Symbol(k) => v for (k, v) in d[:temporal]) + TemporalEmbeddings(; td...) + else + nothing + end + return EmbeddingLayer(num, temp) +end + +# Output width of a temporal embedding: d_embedding, +1 when trend is appended. +temporal_out_dim(t::TemporalEmbeddings) = t.d_embedding + (t.trend ? 1 : 0) + +# Per-feature output width of a numerical embedding. Identity and BatchNorm keep each +# feature at width 1; the expanding embeddings emit `d_embedding` columns per feature. +_num_d_embedding(::IdentityEmbedding) = 1 +_num_d_embedding(::BatchNormEmbeddings) = 1 +_num_d_embedding(s::Union{LinearEmbeddings,PeriodicEmbeddings,PiecewiseLinearEmbeddings}) = s.d_embedding + +""" + needs_x_train(spec) -> Bool - return EmbeddingConfig(embedding_type, d_embedding, activation, bins, frequencies, frequencies_init_scale) +Whether building `spec` requires the training matrix. + +# Arguments +- `spec::AbstractEmbedding`: The embedding spec to inspect. + +# Returns +`true` when `spec` needs `x_train` (e.g. piecewise-linear bin edges or temporal normalization statistics), `false` otherwise. +""" +needs_x_train(::Nothing) = false +needs_x_train(::IdentityEmbedding) = false +needs_x_train(::AbstractNumericalEmbedding) = false +needs_x_train(::PiecewiseLinearEmbeddings) = true +needs_x_train(::AbstractTemporalEmbedding) = true +needs_x_train(e::EmbeddingLayer) = needs_x_train(e.num) || needs_x_train(e.temp) + +# Numerical embeddings: each builds its own chain emitting a flat (width, batch) +# output. Expanding types append their own FlattenLayer; BatchNorm is already 2D and +# Identity is a no-op. The builder never inspects the type to decide how to flatten; +# adding a new numerical embedding means adding one method here. +# Row (feature-column) selector used by the temporal `Parallel` branches. Named (not a +# closure) so the resulting `WrappedFunction` is concretely typed via `Base.Fix2`. +_select_rows(x::AbstractMatrix, rows) = x[rows, :] + +_build_num(::IdentityEmbedding; nfeats::Int, x_train=nothing) = NoOpLayer() +_build_num(s::LinearEmbeddings; nfeats::Int, x_train=nothing) = + Chain(_LinearEmbeddings(; nfeats, d_embedding=s.d_embedding, activation=act_dict[s.activation]), + FlattenLayer()) +_build_num(s::PeriodicEmbeddings; nfeats::Int, x_train=nothing) = + Chain(_PeriodicEmbeddings(; nfeats, d_embedding=s.d_embedding, + frequencies=s.frequencies, frequencies_init_scale=s.frequencies_init_scale, + activation=act_dict[s.activation], lite=s.lite), FlattenLayer()) +function _build_num(s::PiecewiseLinearEmbeddings; nfeats::Int, x_train=nothing) + x_train === nothing && + error("PiecewiseLinearEmbeddings requires x_train to compute bin edges") + Chain(_PiecewiseLinearEmbeddings(; bins=compute_bins(x_train; bins=s.bins), + d_embedding=s.d_embedding, activation=act_dict[s.activation], version=s.version), + FlattenLayer()) +end +_build_num(s::BatchNormEmbeddings; nfeats::Int, x_train=nothing) = _BatchNormEmbeddings(; nfeats) + +function _build_temp(t::TemporalEmbeddings, x_col) + x_col === nothing && + error("TemporalEmbeddings requires x_train for t_mean/t_std") + t_mean = Float32(mean(x_col)) + t_std = length(x_col) > 1 ? max(Float32(std(x_col)), 1f-6) : 1f0 + _TemporalEmbeddings(t_mean, t_std, t.order, t.trend, t.d_embedding; periods=t.periods) end """ - (config::EmbeddingConfig)(; nfeats, x_train=nothing) + build_embedding_chain(spec, nfeats; x_train=nothing) -Build `Chain(embedding, FlattenLayer())` from `config`. +Build the embedding chain for `spec`. A numerical-only spec embeds every column; with a +temporal branch the input is routed through a `Lux.Parallel`, applying `num` to the non-time +columns and `temp` to the time column, concatenated features-first then temporal. # Arguments +- `spec::AbstractEmbedding`: The embedding spec to build. - `nfeats::Int`: Number of input features. -- `x_train`: Training matrix `(n_samples, n_features)`; required for `:piecewise`. -""" -function (config::EmbeddingConfig)(; nfeats::Int, x_train=nothing) - emb = if config.embedding_type == :periodic - PeriodicEmbeddings(nfeats, config.d_embedding; - frequencies=config.frequencies, - frequencies_init_scale=config.frequencies_init_scale, - activation=act_dict[config.activation]) - elseif config.embedding_type == :linear - LinearEmbeddings(nfeats, config.d_embedding; activation=act_dict[config.activation]) - elseif config.embedding_type == :piecewise - @assert x_train !== nothing "Piecewise embeddings require `x_train` to compute bin edges." - bins = compute_bins(x_train; bins=config.bins) - @assert length(bins) == nfeats "Expected $nfeats bin vectors, got $(length(bins))" - PiecewiseLinearEmbeddings(bins, config.d_embedding; activation=act_dict[config.activation]) - elseif config.embedding_type == :batchnorm - BatchNormEmbeddings(nfeats) - else - error("Unsupported embedding type: $(config.embedding_type)") - end - return Chain(emb, FlattenLayer()) +- `x_train`: Training matrix `(n_samples, nfeats)`, required when `needs_x_train(spec)` (default `nothing`). + +# Returns +A `Lux` layer emitting a flat `(width, batch)` output; recover the width with [`embedding_width`](@ref). +""" +build_embedding_chain(s::AbstractNumericalEmbedding, nfeats::Int; x_train=nothing) = + _build_num(s; nfeats, x_train) +function build_embedding_chain(e::EmbeddingLayer, nfeats::Int; x_train=nothing) + # numerical branch only + isnothing(e.temp) && return _build_num(e.num; nfeats, x_train) + + # numerical + temporal: route the non-time columns through `num` and the time + # column through `temp`, concatenating features-first / temporal-last. When `num` + # is IdentityEmbedding this is the temporal-only case (features pass through). + idx = e.temp.index + 1 <= idx <= nfeats || + throw(ArgumentError("temporal index=$idx out of range for nfeats=$nfeats")) + keep = setdiff(1:nfeats, idx) + core = _build_num(e.num; nfeats=nfeats - 1, x_train=x_train === nothing ? nothing : x_train[:, keep]) + temp = _build_temp(e.temp, x_train === nothing ? nothing : @view x_train[:, idx]) + return Parallel(vcat, + Chain(WrappedFunction(Base.Fix2(_select_rows, keep)), core), + Chain(WrappedFunction(Base.Fix2(_select_rows, idx:idx)), temp)) +end + +""" + embedding_width(layer, x, rng) -> Int + +Output width of `layer`, measured by a single forward pass on the probe `x`. A forward pass +is robust to the layer's internal structure, unlike analytic `outputsize` which cannot see +through a `vcat` connection. + +# Arguments +- `layer`: The embedding layer to measure. +- `x::AbstractMatrix`: Probe input of shape `(nfeats, batch)` (`init` passes `(nfeats, 2)`). +- `rng::AbstractRNG`: RNG used for parameter setup. + +# Returns +The flattened output width as an `Int`. +""" +function embedding_width(layer, x, rng::AbstractRNG) + ps, st = LuxCore.setup(rng, layer) + st = LuxCore.testmode(st) + y, _ = layer(x, ps, st) + return size(y, 1) end + +""" + has_real_embedding(spec) -> Bool + +Whether `spec` applies a non-identity embedding. + +# Arguments +- `spec::AbstractEmbedding`: The embedding spec to inspect. + +# Returns +`true` when `spec` has a non-identity `num` or a temporal branch; `false` for a bare `IdentityEmbedding` or an empty `EmbeddingLayer`. +""" +has_real_embedding(::IdentityEmbedding) = false +has_real_embedding(::AbstractNumericalEmbedding) = true +has_real_embedding(e::EmbeddingLayer) = !(e.num isa IdentityEmbedding) || e.temp !== nothing + +""" + per_feature_widths(spec, nfeats) -> Vector{Int} + +Output width contributed by each input feature, summing to the total embedding width. Used by +TabM's grouped scaling init to keep a feature's columns together. Ordering matches the concat +layout: non-time features first (numerical per-feature width), temporal column last. + +# Arguments +- `spec::AbstractEmbedding`: The embedding spec. +- `nfeats::Int`: Number of input features. + +# Returns +A `Vector{Int}` of per-feature widths summing to the total embedding width. +""" +per_feature_widths(::IdentityEmbedding, nfeats::Int) = fill(1, nfeats) +per_feature_widths(s::AbstractNumericalEmbedding, nfeats::Int) = fill(_num_d_embedding(s), nfeats) +function per_feature_widths(e::EmbeddingLayer, nfeats::Int) + d_emb = _num_d_embedding(e.num) + isnothing(e.temp) && return fill(d_emb, nfeats) + return vcat(fill(d_emb, nfeats - 1), Int[temporal_out_dim(e.temp)]) +end \ No newline at end of file diff --git a/src/models/embeddings/embeddings.jl b/src/models/embeddings/embeddings.jl index 918f231..801860f 100644 --- a/src/models/embeddings/embeddings.jl +++ b/src/models/embeddings/embeddings.jl @@ -1,16 +1,16 @@ module Embeddings using Lux -using Lux: BatchNorm, Chain, Dense, FlattenLayer using LuxCore +using Random using NNlib -using Random: AbstractRNG, rand, randn -using Statistics: quantile +import Statistics: mean, quantile, std -export NLinear, LinearEmbeddings -export Periodic, PeriodicEmbeddings -export PiecewiseLinearEncoding, PiecewiseLinearEmbeddings -export compute_bins, EmbeddingConfig +export AbstractNumericalEmbedding, AbstractTemporalEmbedding, AbstractEmbedding +export LinearEmbeddings, PeriodicEmbeddings, PiecewiseLinearEmbeddings +export BatchNormEmbeddings, TemporalEmbeddings, IdentityEmbedding +export EmbeddingLayer, build_embedding_chain, needs_x_train, temporal_out_dim +export per_feature_widths, has_real_embedding, embedding_width include("compute_bins.jl") include("nlinear.jl") @@ -18,6 +18,7 @@ include("linear.jl") include("periodic.jl") include("piecewise_linear.jl") include("batchnorm.jl") +include("temporal.jl") include("config.jl") -end +end \ No newline at end of file diff --git a/src/models/embeddings/linear.jl b/src/models/embeddings/linear.jl index da7d589..32de0cc 100644 --- a/src/models/embeddings/linear.jl +++ b/src/models/embeddings/linear.jl @@ -1,24 +1,27 @@ """ - LinearEmbeddings(nfeats, d_embedding; activation=relu) + _LinearEmbeddings(; nfeats, d_embedding, activation=relu) -Per-feature affine map `activation(w * x + b)`. Output shape `(d_embedding, nfeats, batch)`. +Embeds each continuous feature via a learned affine transformation followed by +an activation: `activation(w_j * x_j + b_j)`. +Produces a `(d_embedding, nfeats, batch)` tensor. # Arguments - `nfeats::Int`: Number of input features. - `d_embedding::Int`: Embedding dimension per feature. -- `activation`: Element-wise activation (default `relu`). +- `activation`: Activation function applied element-wise (default `relu`). + E.g. `relu`, `tanh`, `identity`. """ -struct LinearEmbeddings{F} <: LuxCore.AbstractLuxLayer +struct _LinearEmbeddings{F} <: LuxCore.AbstractLuxLayer nfeats::Int d_embedding::Int activation::F end -function LinearEmbeddings(nfeats::Int, d_embedding::Int; activation=relu) - return LinearEmbeddings(nfeats, d_embedding, activation) +function _LinearEmbeddings(; nfeats::Int, d_embedding::Int, activation=NNlib.relu) + return _LinearEmbeddings(nfeats, d_embedding, activation) end -function LuxCore.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) +function LuxCore.initialparameters(rng::AbstractRNG, l::_LinearEmbeddings) limit = Float32(l.d_embedding)^(-0.5f0) weight = reshape((rand(rng, Float32, l.d_embedding, l.nfeats) .* 2f0 .* limit) .- limit, l.d_embedding, l.nfeats, 1) @@ -27,9 +30,11 @@ function LuxCore.initialparameters(rng::AbstractRNG, l::LinearEmbeddings) return (weight=weight, bias=bias) end -LuxCore.initialstates(::AbstractRNG, ::LinearEmbeddings) = (;) +LuxCore.initialstates(::AbstractRNG, ::_LinearEmbeddings) = (;) -function (l::LinearEmbeddings)(x::AbstractMatrix, ps, st) +function (l::_LinearEmbeddings)(x::AbstractMatrix, ps, st) x_r = reshape(x, 1, size(x, 1), size(x, 2)) return l.activation.(muladd.(ps.weight, x_r, ps.bias)), st end + +LuxCore.outputsize(l::_LinearEmbeddings, x, ::AbstractRNG) = (l.d_embedding, size(x, 1)) \ No newline at end of file diff --git a/src/models/embeddings/periodic.jl b/src/models/embeddings/periodic.jl index a96a144..90c16d4 100644 --- a/src/models/embeddings/periodic.jl +++ b/src/models/embeddings/periodic.jl @@ -1,12 +1,13 @@ """ Periodic(nfeats, n_frequencies, sigma) -Sinusoidal encoding `[cos(2π w x), sin(2π w x)]`. Output shape `(2 * n_frequencies, nfeats, batch)`. +Maps each feature to `2 * n_frequencies` sinusoidal components +`[cos(2π w x), sin(2π w x)]`. Output shape `(2 * n_frequencies, nfeats, batch)`. # Arguments - `nfeats::Int`: Number of input features. -- `n_frequencies::Int`: Frequency components per feature. -- `sigma::Float32`: Std-dev for frequency weight init (clamped to ±3σ). +- `n_frequencies::Int`: Number of frequency components per feature. +- `sigma::Float32`: Std-dev for the frequency weight initialization (clamped to ±3σ). """ struct Periodic <: LuxCore.AbstractLuxLayer nfeats::Int @@ -30,32 +31,34 @@ function (l::Periodic)(x::AbstractMatrix, ps, st) end """ - PeriodicEmbeddings(nfeats, d_embedding=24; frequencies=48, - frequencies_init_scale=0.01f0, activation=relu, lite=false) + _PeriodicEmbeddings(; nfeats, d_embedding=24, frequencies=48, + frequencies_init_scale=0.01f0, activation=relu, lite=false) -`Periodic` followed by per-feature linear projection and activation. +Periodic sinusoidal encoding followed by a learned linear projection: applies +`Periodic`, then `NLinear` (or a shared `Dense` if `lite`), then `activation`. # Arguments - `nfeats::Int`: Number of input features. -- `d_embedding::Int`: Output dimension per feature (default `24`). -- `frequencies::Int`: Sinusoidal components per feature (default `48`). -- `frequencies_init_scale::Float32`: σ for frequency init (default `0.01f0`). -- `activation`: Post-projection activation (default `relu`). -- `lite::Bool`: Use shared `Dense` instead of per-feature `NLinear` (default `false`). +- `d_embedding::Int`: Output embedding dimension per feature (default `24`). +- `frequencies::Int`: Sinusoidal frequency components per feature (default `48`). +- `frequencies_init_scale::Float32`: σ for frequency weight init (default `0.01f0`). +- `activation`: Activation applied after projection (default `relu`). +- `lite::Bool`: Use a single shared `Dense` instead of per-feature `NLinear` + (default `false`). Requires a non-identity `activation`. """ -struct PeriodicEmbeddings{P,L,F} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :linear)} +struct _PeriodicEmbeddings{P,L,F} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :linear)} periodic::P linear::L activation::F lite::Bool end -function PeriodicEmbeddings( +function _PeriodicEmbeddings(; nfeats::Int, - d_embedding::Int=24; + d_embedding::Int=24, frequencies::Int=48, frequencies_init_scale::Float32=0.01f0, - activation=relu, + activation=NNlib.relu, lite::Bool=false, ) if lite && activation === identity @@ -67,10 +70,10 @@ function PeriodicEmbeddings( else NLinear(nfeats, 2 * frequencies, d_embedding) end - return PeriodicEmbeddings(periodic, linear, activation, lite) + return _PeriodicEmbeddings(periodic, linear, activation, lite) end -function (m::PeriodicEmbeddings)(x::AbstractMatrix, ps, st) +function (m::_PeriodicEmbeddings)(x::AbstractMatrix, ps, st) h, st_p = m.periodic(x, ps.periodic, st.periodic) h_lin, st_l = if m.lite @@ -82,5 +85,10 @@ function (m::PeriodicEmbeddings)(x::AbstractMatrix, ps, st) m.linear(h, ps.linear, st.linear) end - return m.activation.(h_lin), (periodic=st_p, linear=st_l) + h_lin = m.activation.(h_lin) + + return h_lin, (periodic=st_p, linear=st_l) end + +LuxCore.outputsize(m::_PeriodicEmbeddings, x, ::AbstractRNG) = + m.lite ? error("lite _PeriodicEmbeddings outputsize undefined") : (m.linear.out_features, size(x, 1)) \ No newline at end of file diff --git a/src/models/embeddings/piecewise_linear.jl b/src/models/embeddings/piecewise_linear.jl index 223a7a4..bfa332f 100644 --- a/src/models/embeddings/piecewise_linear.jl +++ b/src/models/embeddings/piecewise_linear.jl @@ -1,7 +1,7 @@ """ PiecewiseLinearEncoding(bins) -Non-trainable piecewise-linear encoding from precomputed bin edges. +Non-trainable piecewise-linear encoding using precomputed bin edges. Output shape `(max_n_bins, nfeats, batch)`. # Arguments @@ -36,6 +36,9 @@ function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) b = -bin_edges[1:end-1] ./ bin_width nb = length(bin_edges) - 1 + # Place the last bin's weight/bias at the end row; + # remaining bins fill rows 1:nb-1. Unused rows stay zero + # and are clamped to [0, 1] harmlessly. weight[end, i] = w[end] bias[end, i] = b[end] if nb > 1 @@ -44,6 +47,7 @@ function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) end end + # Pre-reshape to 3D to avoid per-call reshape in forward pass return ( weight=reshape(weight, M, N, 1), bias=reshape(bias, M, N, 1), @@ -51,25 +55,27 @@ function LuxCore.initialstates(::AbstractRNG, l::PiecewiseLinearEncoding) end function (l::PiecewiseLinearEncoding)(x::AbstractMatrix, ps, st) + # For each feature j and bin b: h[b,j,:] = clamp(w[b,j]*x[j,:] + bias[b,j], 0, 1) x_r = reshape(x, 1, size(x, 1), size(x, 2)) h = clamp.(muladd.(st.weight, x_r, st.bias), 0f0, 1f0) return h, st end """ - PiecewiseLinearEmbeddings(bins, d_embedding; activation=identity, version=:B) + _PiecewiseLinearEmbeddings(; bins, d_embedding, activation=identity, version=:B) -Learnable projection on top of [`PiecewiseLinearEncoding`](@ref). -Version `:A`: encoding → `NLinear` (with bias). -Version `:B`: encoding → zero-init `NLinear` + [`LinearEmbeddings`](@ref) residual. +Learnable embeddings on top of `PiecewiseLinearEncoding`. +Version `:A`: PLE -> NLinear (with bias). +Version `:B`: PLE -> NLinear (zero-init, no bias) + per-feature linear residual. +Output shape `(d_embedding, nfeats, batch)`. # Arguments -- `bins::Vector{<:AbstractVector}`: Bin edges from [`compute_bins`](@ref). -- `d_embedding::Int`: Output dimension per feature. -- `activation`: Post-projection activation (default `identity`). +- `bins::Vector{<:AbstractVector}`: Bin edges per feature from [`compute_bins`](@ref). +- `d_embedding::Int`: Embedding dimension per feature. +- `activation`: Activation function applied after projection (default `identity`). E.g. `relu`, `tanh`. - `version::Symbol`: `:A` or `:B` (default `:B`). """ -struct PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)} +struct _PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{(:linear0, :encoding, :linear)} linear0::L0 encoding::I linear::L @@ -77,9 +83,9 @@ struct PiecewiseLinearEmbeddings{L0,I,L,F} <: LuxCore.AbstractLuxContainerLayer{ version::Symbol end -function PiecewiseLinearEmbeddings( +function _PiecewiseLinearEmbeddings(; bins::Vector{<:AbstractVector}, - d_embedding::Int; + d_embedding::Int, activation=identity, version::Symbol=:B, ) @@ -88,13 +94,14 @@ function PiecewiseLinearEmbeddings( max_n_bins = maximum(length(b) - 1 for b in bins) encoding = PiecewiseLinearEncoding(bins) - linear0 = (version == :B) ? LinearEmbeddings(nfeats, d_embedding; activation=identity) : nothing + # Residual path uses raw affine output (no activation) + linear0 = (version == :B) ? _LinearEmbeddings(; nfeats, d_embedding, activation=identity) : nothing linear = NLinear(nfeats, max_n_bins, d_embedding; bias=(version == :A)) - return PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version) + return _PiecewiseLinearEmbeddings(linear0, encoding, linear, activation, version) end -function LuxCore.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbeddings) +function LuxCore.initialparameters(rng::AbstractRNG, m::_PiecewiseLinearEmbeddings) ps_l0 = m.linear0 === nothing ? nothing : LuxCore.initialparameters(rng, m.linear0) ps_enc = LuxCore.initialparameters(rng, m.encoding) @@ -108,7 +115,7 @@ function LuxCore.initialparameters(rng::AbstractRNG, m::PiecewiseLinearEmbedding return (linear0=ps_l0, encoding=ps_enc, linear=ps_lin) end -function (m::PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st) +function (m::_PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st) val_linear0 = nothing st_l0 = st.linear0 @@ -124,3 +131,5 @@ function (m::PiecewiseLinearEmbeddings)(x::AbstractMatrix, ps, st) return h_final, (linear0=st_l0, encoding=st_enc, linear=st_lin) end + +LuxCore.outputsize(m::_PiecewiseLinearEmbeddings, x, ::AbstractRNG) = (m.linear.out_features, size(x, 1)) \ No newline at end of file diff --git a/src/models/embeddings/temporal.jl b/src/models/embeddings/temporal.jl new file mode 100644 index 0000000..fd37433 --- /dev/null +++ b/src/models/embeddings/temporal.jl @@ -0,0 +1,77 @@ +""" +Default Fourier base periods for `TemporalEmbeddings`, in seconds: +year (`365.25 * 86_400`), 30-day month, week, day. Assumes the time column is +encoded as a POSIX-style seconds-since-epoch float. +""" +const _DEFAULT_TEMPORAL_PERIODS = Float32[31_557_600, 2_629_800, 604_800, 86_400] + +"""Fixed Fourier features `sin(ωᵢ t), cos(ωᵢ t)` per harmonic; `ωᵢ = 2πk/T`.""" +struct TemporalPeriodic <: LuxCore.AbstractLuxLayer + omega::Vector{Float32} +end + +LuxCore.initialparameters(::AbstractRNG, ::TemporalPeriodic) = (;) + +# `omega` is fixed (no gradient), but it lives in `st` rather than as a struct +# field so that `dev(st)` transfers it to the GPU alongside the learnable +# state. Reshaped to a column so the broadcast against `(1, batch)` x is unambiguous. +LuxCore.initialstates(::AbstractRNG, l::TemporalPeriodic) = + (omega = Matrix{Float32}(reshape(l.omega, length(l.omega), 1)),) + +function (l::TemporalPeriodic)(x::AbstractMatrix, ps, st) + z = st.omega .* x + return vcat(sin.(z), cos.(z)), st +end + +"""Realized temporal embedding: Fourier features + dense projection (+ optional linear trend).""" +struct _TemporalEmbeddings{P,D,trend} <: LuxCore.AbstractLuxContainerLayer{(:periodic, :dense)} + periodic::P + dense::D + t_mean::Float32 + t_std::Float32 +end + +""" + _TemporalEmbeddings(t_mean, t_std, order, trend, d_embedding; periods) + +Realize a `_TemporalEmbeddings` layer from spec. + +`order` and `periods` are aligned per-band: for each `(o_i, p_i)`, the Fourier +basis contributes harmonics `k = 1:o_i` with angular frequency `ω = 2πk/p_i`. +All `ω` are concatenated into a single flat vector, so the resulting +`TemporalPeriodic` outputs `2 * sum(order)` features (sin + cos per harmonic), +which the `Dense` then projects to `d_embedding`. + +`trend` is lifted into the type parameter so the forward dispatches statically +on whether the standardised raw time `(x - t_mean) / t_std` is appended. +""" +function _TemporalEmbeddings( + t_mean::Real, t_std::Real, + order::AbstractVector{<:Integer}, trend::Bool, d_embedding::Int; + periods::AbstractVector{<:Real}=_DEFAULT_TEMPORAL_PERIODS, +) + @assert length(order) == length(periods) "length(order) must equal length(periods)" + @assert any(>(0), order) "`order` must contain at least one positive entry" + omega = Float32[2f0 * Float32(π) * Float32(k) / Float32(p) + for (o, p) in zip(order, periods) for k in 1:o] + periodic = TemporalPeriodic(omega) + dense = Dense(2 * length(omega) => d_embedding, NNlib.relu) + return _TemporalEmbeddings{typeof(periodic),typeof(dense),trend}( + periodic, dense, Float32(t_mean), Float32(t_std), + ) +end + +function (m::_TemporalEmbeddings{P,D,true})(x::AbstractMatrix, ps, st) where {P,D} + h, st_p = m.periodic(x, ps.periodic, st.periodic) + out, st_d = m.dense(h, ps.dense, st.dense) + return vcat(out, (x .- m.t_mean) ./ m.t_std), (periodic=st_p, dense=st_d) +end + +function (m::_TemporalEmbeddings{P,D,false})(x::AbstractMatrix, ps, st) where {P,D} + h, st_p = m.periodic(x, ps.periodic, st.periodic) + out, st_d = m.dense(h, ps.dense, st.dense) + return out, (periodic=st_p, dense=st_d) +end + +LuxCore.outputsize(m::_TemporalEmbeddings{P,D,true}, x, ::AbstractRNG) where {P,D} = (m.dense.out_dims + 1,) +LuxCore.outputsize(m::_TemporalEmbeddings{P,D,false}, x, ::AbstractRNG) where {P,D} = (m.dense.out_dims,) \ No newline at end of file diff --git a/src/models/models.jl b/src/models/models.jl index 3e8e3d1..73cbf5e 100644 --- a/src/models/models.jl +++ b/src/models/models.jl @@ -1,8 +1,11 @@ module Models export NeuroTabModel, Architecture -export Embeddings -export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig, MOETreeConfig +export Embeddings, EmbeddingLayer +export LinearEmbeddings, PeriodicEmbeddings, PiecewiseLinearEmbeddings +export BatchNormEmbeddings, TemporalEmbeddings, IdentityEmbedding +export AbstractNumericalEmbedding, AbstractTemporalEmbedding, AbstractEmbedding +export NeuroTreeConfig, MLPConfig, ResNetConfig, TabMConfig, MOETreeConfig, ModernNCAConfig using ..Losses using Lux: Chain @@ -32,16 +35,60 @@ function get_activation(act::Symbol) return activation_dict[act] end +""" + train_dataloader(arch, m, default, df; kwargs...) + +Per-architecture hook: return the dataloader `fit` should use. Default returns +the tabular `default` unchanged; retrieval-style archs override to substitute +their own iterator (e.g. with a candidate corpus attached). + +Overrides may write into `m.info` (e.g. to stash a corpus for inference). +""" +train_dataloader(::Architecture, ::Any, data, ::Any; kw...) = data + +""" + build_chain(arch, embed_chain; nfeats, outsize, d_in, kw...) + +Per-architecture hook for assembling the Lux chain that `fit` will train. +The embedding chain always exists (identity is `WrappedFunction(identity)`); +`d_in` is supplied by the caller. Default wires `embed_chain` in front of the +architecture's backbone via `Chain(embed_chain, arch(...))`. +""" +build_chain(arch::Architecture, embed_chain; nfeats, outsize, d_in, kw...) = + Chain(embed_chain, arch(; nfeats=d_in, outsize)) + +""" + infer_dataloader(chain, info, data, dev, ps, st) + +Per-architecture hook: return the per-batch iterator `infer` should use. +Default returns `data` unchanged; retrieval-style archs override to wrap each +batch with extra inputs (e.g. a candidate corpus). Dispatched on +`typeof(m.chain)` so arch modules can override on their concrete model type. +""" +infer_dataloader(::Any, ::Any, data, ::Any, ::Any, ::Any) = data + +""" + eval_dataloader(chain, info, data, dev, ps, st) + +Per-architecture hook: return the per-batch iterator the eval `CallBack` should +use. Default returns `data` unchanged; retrieval-style archs override to wrap +each batch's `x` with extra inputs (e.g. a candidate corpus) so that the eval +forward pass matches the training/inference signature. Dispatched on +`typeof(m.chain)`. +""" +eval_dataloader(::Any, ::Any, data, ::Any, ::Any, ::Any) = data + """ NeuroTabModel The object containing the model and associated metadata. -## Fields +# Fields -- `loss_type`: the loss function type used during training (e.g. `MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`) -- `chain`: the underlying `Lux.Chain` neural network -- `info`: a `Dict{Symbol,Any}` storing metadata such as `:feature_names`, `:target_levels`, `:device`, `logger`, as well as fitted parameters (`ps`) and state (`st`). +- `loss_type`: the loss type used in training (`MSE`, `LogLoss`, `MLogLoss`, `GaussianMLE`). +- `chain`: the underlying `Lux.Chain` neural network. +- `info`: a `Dict{Symbol,Any}` of metadata such as `:feature_names`, `:target_levels`, + and `:device`, plus the fitted parameters (`ps`) and state (`st`). """ struct NeuroTabModel{L<:LossType,C} loss_type::Type{L} @@ -61,10 +108,19 @@ using .MOETrees include("TabM/TabM.jl") using .TabM +function build_chain(arch::TabMConfig, embed_chain; nfeats, outsize, d_in, embedding_config, kw...) + d_features = per_feature_widths(embedding_config, nfeats) + scaling_init_override = has_real_embedding(embedding_config) ? :normal : nothing + return Chain(embed_chain, arch(; nfeats=d_in, outsize, d_features, scaling_init_override)) +end + include("MLP/mlp.jl") using .MLP include("ResNet/resnet.jl") using .ResNet +include("ModernNCA/modernnca.jl") +using .ModernNCA + end \ No newline at end of file