diff --git a/docs/src/api.md b/docs/src/api.md index ffc7401..8834036 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -86,6 +86,11 @@ ReLUQuadratic ReLUSOS1 ``` +## `Scale` +```@docs +Scale +``` + ## `Sigmoid` ```@docs Sigmoid diff --git a/docs/src/manual/predictors.md b/docs/src/manual/predictors.md index 1df23f1..6f0553e 100644 --- a/docs/src/manual/predictors.md +++ b/docs/src/manual/predictors.md @@ -20,6 +20,7 @@ The following predictors are supported. See their docstrings for details: | [`ReLUBigM`](@ref) | $f(x) = \max.(0, x)$ | $M \rightarrow M$ | | [`ReLUQuadratic`](@ref) | $f(x) = \max.(0, x)$ | $M \rightarrow M$ | | [`ReLUSOS1`](@ref) | $f(x) = \max.(0, x)$ | $M \rightarrow M$ | +| [`Scale`](@ref) | $f(x) = scale .* x .+ bias$ | $M \rightarrow M$ | | [`Sigmoid`](@ref) | $f(x) = \frac{1}{1 + e^{-x}}$ | $M \rightarrow M$ | | [`SoftMax`](@ref) | $f(x) = \frac{e^{x_i}}{\sum e^{x_i}}$ | $M \rightarrow 1$ | | [`SoftPlus`](@ref) | $f(x) = \log(1 + e^x)$ | $M \rightarrow M$ | diff --git a/ext/MathOptAIFluxExt.jl b/ext/MathOptAIFluxExt.jl index 4c2945b..c3df772 100644 --- a/ext/MathOptAIFluxExt.jl +++ b/ext/MathOptAIFluxExt.jl @@ -24,6 +24,7 @@ Add a trained neural network from Flux.jl to `model`. ## Supported layers * `Flux.Dense` + * `Flux.Scale` * `Flux.softmax` ## Supported activation functions @@ -86,6 +87,7 @@ Convert a trained neural network from Flux.jl to a [`Pipeline`](@ref). ## Supported layers * `Flux.Dense` + * `Flux.Scale` * `Flux.softmax` ## Supported activation functions @@ -173,4 +175,14 @@ function _add_predictor( return end +function _add_predictor( + predictor::MathOptAI.Pipeline, + layer::Flux.Scale, + config::Dict, +) + push!(predictor.layers, MathOptAI.Scale(layer.scale, layer.bias)) + _add_predictor(predictor, layer.σ, config) + return +end + end # module diff --git a/ext/MathOptAILuxExt.jl b/ext/MathOptAILuxExt.jl index 24ed369..57dc99f 100644 --- a/ext/MathOptAILuxExt.jl +++ b/ext/MathOptAILuxExt.jl @@ -24,6 +24,7 @@ Add a trained neural network from Lux.jl to `model`. ## Supported layers * `Lux.Dense` + * `Lux.Scale` ## Supported activation functions @@ -96,6 +97,7 @@ Convert a trained neural network from Lux.jl to a [`Pipeline`](@ref). ## Supported layers * `Lux.Dense` + * `Lux.Scale` ## Supported activation functions @@ -188,4 +190,15 @@ function _add_predictor( return end +function _add_predictor( + predictor::MathOptAI.Pipeline, + layer::Lux.Scale, + p, + config::Dict, +) + push!(predictor.layers, MathOptAI.Scale(p.weight, p.bias)) + _add_predictor(predictor, layer.activation, config) + return +end + end # module diff --git a/src/predictors/Scale.jl b/src/predictors/Scale.jl new file mode 100644 index 0000000..7910573 --- /dev/null +++ b/src/predictors/Scale.jl @@ -0,0 +1,82 @@ +# Copyright (c) 2024: Oscar Dowson and contributors +# Copyright (c) 2024: Triad National Security, LLC +# +# Use of this source code is governed by a BSD-style license that can be found +# in the LICENSE.md file. + +""" + Scale( + scale::Vector{T}, + bias::Vector{T}, + ) where {T} <: AbstractPredictor + +An [`AbstractPredictor`](@ref) that represents the affine relationship: +```math +f(x) = scale .* x .+ bias +``` + +## Example + +```jldoctest +julia> using JuMP, MathOptAI + +julia> model = Model(); + +julia> @variable(model, x[1:2]); + +julia> f = MathOptAI.Scale([2.0, 3.0], [4.0, 5.0]) +Scale(scale, bias) + +julia> y = MathOptAI.add_predictor(model, f, x) +1-element Vector{VariableRef}: + moai_Scale[1] + moai_Scale[2] + +julia> print(model) +Feasibility +Subject to + 2 x[1] + moai_Scale[1] = -4 + 3 x[2] + moai_Scale[2] = -5 + +julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x) +2-element Vector{AffExpr}: + 2 x[1] + 4 + 3 x[2] + 5 +``` +""" +struct Scale{T} <: AbstractPredictor + scale::Vector{T} + bias::Vector{T} +end + +function Base.show(io::IO, ::Scale) + return print(io, "Scale(scale, bias)") +end + +function add_predictor( + model::JuMP.AbstractModel, + predictor::Scale, + x::Vector, +) + m = length(predictor.scale) + y = JuMP.@variable(model, [1:m], base_name = "moai_Scale") + bounds = _get_variable_bounds.(x) + for (i, scale) in enumerate(predictor.scale) + y_lb = y_ub = predictor.bias[i] + lb, ub = bounds[i] + y_ub += scale * ifelse(scale >= 0, ub, lb) + y_lb += scale * ifelse(scale >= 0, lb, ub) + _set_bounds_if_finite(y[i], y_lb, y_ub) + end + JuMP.@constraint(model, predictor.scale .* x .+ predictor.bias .== y) + return y +end + +function add_predictor( + model::JuMP.AbstractModel, + predictor::ReducedSpace{<:Scale}, + x::Vector, +) + scale, bias = predictor.predictor.scale, predictor.predictor.bias + return JuMP.@expression(model, scale .* x .+ bias) +end diff --git a/test/test_Flux.jl b/test/test_Flux.jl index 2808095..2b9c3d4 100644 --- a/test/test_Flux.jl +++ b/test/test_Flux.jl @@ -41,6 +41,31 @@ function _train_lux_model(model) return model end +function test_end_to_end_with_scale() + chain = _train_lux_model( + Flux.Chain( + Flux.Scale(1), + Flux.Dense(1 => 16, Flux.relu), + Flux.Dense(16 => 1), + ), + ) + model = Model(HiGHS.Optimizer) + set_silent(model) + @variable(model, x) + y = MathOptAI.add_predictor( + model, + chain, + [x]; + config = Dict(Flux.relu => MathOptAI.ReLUBigM(100.0)), + ) + @constraint(model, only(y) <= 4) + @objective(model, Min, x) + optimize!(model) + @test is_solved_and_feasible(model) + @test isapprox(value(x), -1.24; atol = 1e-1) + return +end + function test_end_to_end_ReLUBigM() chain = _train_lux_model( Flux.Chain(Flux.Dense(1 => 16, Flux.relu), Flux.Dense(16 => 1)), diff --git a/test/test_Lux.jl b/test/test_Lux.jl index a26d114..42b79cf 100644 --- a/test/test_Lux.jl +++ b/test/test_Lux.jl @@ -49,6 +49,31 @@ function _train_lux_model(model) return (model, parameters, state) end +function test_end_to_end_with_scale() + state = _train_lux_model( + Lux.Chain( + Lux.Scale(1), + Lux.Dense(1 => 16, Lux.relu), + Lux.Dense(16 => 1), + ), + ) + model = Model(HiGHS.Optimizer) + set_silent(model) + @variable(model, x) + y = MathOptAI.add_predictor( + model, + state, + [x]; + config = Dict(Lux.relu => MathOptAI.ReLUBigM(100.0)), + ) + @constraint(model, only(y) <= 4) + @objective(model, Min, x) + optimize!(model) + @test is_solved_and_feasible(model) + @test isapprox(value(x), -1.24; atol = 1e-2) + return +end + function test_end_to_end_ReLUBigM() state = _train_lux_model( Lux.Chain(Lux.Dense(1 => 16, Lux.relu), Lux.Dense(16 => 1)), diff --git a/test/test_predictors.jl b/test/test_predictors.jl index 9a6799c..083d5af 100644 --- a/test/test_predictors.jl +++ b/test/test_predictors.jl @@ -337,6 +337,26 @@ function test_ReducedSpace_ReducedSpace() return end +function test_Scale() + model = Model() + @variable(model, x[1:2]) + f = MathOptAI.Scale([2.0, 3.0], [4.0, 5.0]) + @test sprint(show, f) == "Scale(scale, bias)" + y = MathOptAI.add_predictor(model, f, x) + cons = all_constraints(model; include_variable_in_set_constraints = false) + @test length(cons) == 2 + objs = constraint_object.(cons) + @test objs[1].set == MOI.EqualTo(-4.0) + @test objs[2].set == MOI.EqualTo(-5.0) + @test isequal_canonical(objs[1].func, 2.0 * x[1] - y[1]) + @test isequal_canonical(objs[2].func, 3.0 * x[2] - y[2]) + y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x) + cons = all_constraints(model; include_variable_in_set_constraints = false) + @test length(cons) == 2 + @test isequal_canonical(y, [2.0 * x[1] + 4.0, 3.0 * x[2] + 5.0]) + return +end + end # module TestPredictors.runtests()