diff --git a/Project.toml b/Project.toml index 350eabe..934d704 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "MLJLinearModels" uuid = "6ee0df7b-362f-4a72-a706-9e79364fb692" -version = "0.10.3" +version = "0.10.4" authors = ["Thibaut Lienart "] [deps] diff --git a/src/fit/proxgrad.jl b/src/fit/proxgrad.jl index 1a04fe1..a57d34d 100644 --- a/src/fit/proxgrad.jl +++ b/src/fit/proxgrad.jl @@ -21,6 +21,8 @@ function _fit(glr::GLR, solver::ProxGrad, X, y, scratch) # functions _f = if solver.gram smooth_gram_objective(glr, X, y, n) + elseif c == 0 + smooth_objective!(glr, X, y, scratch) else smooth_objective(glr, X, y; c=c) end diff --git a/src/glr/d_l2loss.jl b/src/glr/d_l2loss.jl index 1184c9f..e7bb318 100644 --- a/src/glr/d_l2loss.jl +++ b/src/glr/d_l2loss.jl @@ -69,10 +69,20 @@ function smooth_fg!(glr::GLR{L2Loss,<:ENR}, X, y, scratch) apply_Xt!(g, X, r) g .+= λ .* θ glr.fit_intercept && (glr.penalize_intercept || (g[end] -= λ * θ[end])) - return glr.loss(r) + get_l2(glr.penalty)(view_θ(glr, θ)) + return l2loss!(r) + get_l2(glr.penalty)(view_θ(glr, θ)) end end +# residuals (Xθ-y) formed in scratch.n2, |y-Xθ|₂² = |Xθ-y|₂² +function smooth_objective!(glr::GLR{L2Loss,<:ENR}, X, y, scratch) + J = smooth_objective(glr, size(X, 1)) + r = scratch.n2 + return θ -> (get_residuals!(r, X, θ, y); l2loss!(r) + J.penalty(view_θ(glr, θ))) +end + +# L2Loss()(r) overwriting r; sums the squares as lp(r, 2) does so values match +l2loss!(r) = sum(r .= abs2.(r)) / 2 + function smooth_gram_fg!(glr::GLR{L2Loss,<:ENR}, XX, Xy, n) λ = get_penalty_scale_l2(glr, n) (g, θ) -> begin diff --git a/src/glr/utils.jl b/src/glr/utils.jl index 102da1d..1fb73e8 100644 --- a/src/glr/utils.jl +++ b/src/glr/utils.jl @@ -33,6 +33,19 @@ smooth_objective(glr::GLR, X, y; c::Int=0) = """ $SIGNATURES +Return a function computing the smooth part of the objective at a given +evaluation point `θ`, as `smooth_objective(glr, X, y)` for a single-output +model but writing `X*θ` in `scratch.n2`. +""" +function smooth_objective!(glr::GLR, X, y, scratch) + J = smooth_objective(glr, size(X, 1)) + Xθ = scratch.n2 + return θ -> (apply_X!(Xθ, X, θ); J(y, Xθ, view_θ(glr, θ))) +end + +""" +$SIGNATURES + Return the smooth part of the objective function of a GLR. """ smooth_objective(glr::GLR{<:SmoothLoss,<:ENR}, n) = glr.loss + get_l2(glr.penalty) * ifelse(glr.scale_penalty_with_samples, n, 1.) diff --git a/test/fit/ols-ridge-lasso-elnet.jl b/test/fit/ols-ridge-lasso-elnet.jl index 3c9a1c6..c3429fb 100644 --- a/test/fit/ols-ridge-lasso-elnet.jl +++ b/test/fit/ols-ridge-lasso-elnet.jl @@ -161,3 +161,40 @@ end @test isapprox(θ_fista, θ_gram_explicit, rtol=1e-5) @test isapprox(θ_gram_explicit, θ_gram_implicit; rtol=1e-5) end + +@testset "proxgrad in-place objective" begin + n, p = 20_000, 10 + ((X, _, _), (_, y1, _)) = generate_continuous(n, p; seed=5) + yb = sign.(y1) + # (model, target, check allocations); scaled elnet does not converge here, + # see https://github.com/JuliaAI/MLJLinearModels.jl/issues/187 + for (glr, t, check_alloc) in ( + (LassoRegression(0.5), y1, true), + (LassoRegression(0.5; fit_intercept=false), y1, true), + (ElasticNetRegression(0.3, 0.5), y1, false), + (ElasticNetRegression(0.3, 0.5; scale_penalty_with_samples=false), y1, true), + (LogisticRegression(0.0, 0.2; penalty=:l1), yb, false), + ) + s = R.scratch(X; i=glr.fit_intercept) + f = R.smooth_objective(glr, X, t) + f! = R.smooth_objective!(glr, X, t, s) + r = StableRNG(1) + for _ in 1:5 + θ = randn(r, p + Int(glr.fit_intercept)) + @test f!(θ) == f(θ) + end + check_alloc || continue + # no length-n allocation in an objective or gradient call + fg! = R.smooth_fg!(glr, X, t, s) + θ = randn(r, p + Int(glr.fit_intercept)) + g = similar(θ) + f!(θ); fg!(g, θ) + @test (@allocated f!(θ)) < 8n + @test (@allocated fg!(g, θ)) < 8n + # a fit allocates the 3 length-n scratch vectors and nothing else of length n + for solver in (FISTA(), ISTA()) + fit(glr, X, t; solver=solver) + @test (@allocated fit(glr, X, t; solver=solver)) < 4 * 8n + end + end +end