diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index afb9801..0aacbd4 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -18,7 +18,6 @@ jobs: fail-fast: false matrix: version: - - '1.6' - '1.9' os: - ubuntu-latest diff --git a/Project.toml b/Project.toml index 5ce1825..69e4768 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "L2O" uuid = "e1d8bfa7-c465-446a-84b9-451470f6e76c" authors = ["andrewrosemberg and contributors"] -version = "1.2.0-DEV" +version = "1.0.0" [deps] Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45" @@ -21,12 +21,12 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] -Arrow = "2" -CSV = "0.10" -JuMP = "1" -ParametricOptInterface = "0.7" +Arrow = "^2" +CSV = "^0.10" +JuMP = "^1" +ParametricOptInterface = "^0.8" Zygote = "^0.6.68" -julia = "1.6" +julia = "^1.9" [extras] CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" diff --git a/src/FullyConnected.jl b/src/FullyConnected.jl index 19b7042..67fbef6 100644 --- a/src/FullyConnected.jl +++ b/src/FullyConnected.jl @@ -66,14 +66,20 @@ end # @forward((ConvexRegressor, :model), MLJFlux.Regressor) # Define a container to hold any optimiser specific parameters (if any): -struct ConvexRule <: Flux.Optimise.AbstractOptimiser - rule::Flux.Optimise.AbstractOptimiser +struct ConvexRule <: Optimisers.AbstractRule + rule::Optimisers.AbstractRule tol::Real end -function ConvexRule(rule::Flux.Optimise.AbstractOptimiser; tol=1e-6) +function ConvexRule(rule::Optimisers.AbstractRule; tol=1e-6) return ConvexRule(rule, tol) end +Optimisers.init(o::ConvexRule, x::AbstractArray) = Optimisers.init(o.rule, x) + +function Optimisers.apply!(o::ConvexRule, mvel, x::AbstractArray{T}, dx) where T + return Optimisers.apply!(o.rule, mvel, x, dx) +end + """ function make_convex!(chain::PairwiseFusion; tol = 1e-6) @@ -102,24 +108,48 @@ function make_convex!(model::Chain; tol=1e-6) end end -function MLJFlux.train!( - model::MLJFlux.MLJFluxDeterministic, penalty, chain, optimiser::ConvexRule, X, y -) +function MLJFlux.train( + model, + chain, + optimiser::ConvexRule, + optimiser_state, + epochs, + verbosity, + X, + y, + ) + loss = model.loss + + # intitialize and start progress meter: + meter = MLJFlux.Progress(epochs + 1, dt=0, desc="Optimising neural net:", + barglyphs=MLJFlux.BarGlyphs("[=> ]"), barlen=25, color=:yellow) + verbosity != 1 || MLJFlux.next!(meter) + + # initiate history: n_batches = length(y) - training_loss = zero(Float32) - for i in 1:n_batches - parameters = Flux.params(chain) - gs = Flux.gradient(parameters) do - yhat = chain(X[i]) - batch_loss = loss(yhat, y[i]) + penalty(parameters) / n_batches - training_loss += batch_loss - return batch_loss - end - Flux.update!(optimiser.rule, parameters, gs) + + losses = (loss(chain(X[i]), y[i]) for i in 1:n_batches) + history = [mean(losses),] + + for i in 1:epochs + chain, optimiser_state, current_loss = MLJFlux.train_epoch( + model, + chain, + optimiser, + optimiser_state, + X, + y, + ) make_convex!(chain; tol=optimiser.tol) + verbosity < 2 || + @info "Loss is $(round(current_loss; sigdigits=4))" + verbosity != 1 || next!(meter) + push!(history, current_loss) end - return training_loss / n_batches + + return chain, optimiser_state, history + end function train!(model, loss, opt_state, X, Y; _batchsize=32, shuffle=true) diff --git a/src/datasetgen.jl b/src/datasetgen.jl index 022b20c..610adc0 100644 --- a/src/datasetgen.jl +++ b/src/datasetgen.jl @@ -105,6 +105,14 @@ end abstract type AbstractProblemIterator end +abstract type AbstractParameterType end + +abstract type POIParamaterType <: AbstractParameterType end + +abstract type JuMPNLPParameterType <: AbstractParameterType end + +abstract type JuMPParameterType <: AbstractParameterType end + """ ProblemIterator(ids::Vector{UUID}, pairs::Dict{VariableRef, Vector{Real}}) @@ -115,24 +123,30 @@ struct ProblemIterator{T<:Real} <: AbstractProblemIterator ids::Vector{UUID} pairs::Dict{VariableRef,Vector{T}} early_stop::Function + param_type::Type{<:AbstractParameterType} + pre_solve_hook::Function function ProblemIterator( ids::Vector{UUID}, pairs::Dict{VariableRef,Vector{T}}, early_stop::Function=(args...) -> false, + param_type::Type{<:AbstractParameterType}=POIParamaterType, + pre_solve_hook::Function=(args...) -> nothing ) where {T<:Real} model = JuMP.owner_model(first(keys(pairs))) for (p, val) in pairs @assert length(ids) == length(val) end - return new{T}(model, ids, pairs, early_stop) + return new{T}(model, ids, pairs, early_stop, param_type, pre_solve_hook) end end function ProblemIterator( - pairs::Dict{VariableRef,Vector{T}}; early_stop::Function=(args...) -> false -) where {T<:Real} + pairs::Dict{VariableRef,Vector{T}}; early_stop::Function=(args...) -> false, + pre_solve_hook::Function=(args...) -> nothing, + param_type::Type{<:AbstractParameterType}=POIParamaterType, ids = [uuid1() for _ in 1:length(first(values(pairs)))] - return ProblemIterator(ids, pairs, early_stop) +) where {T<:Real} + return ProblemIterator(ids, pairs, early_stop, param_type, pre_solve_hook) end """ @@ -174,7 +188,8 @@ end function load(model_file::AbstractString, input_file::AbstractString, ::Type{T}; batch_size::Union{Nothing, Integer}=nothing, - ignore_ids::Vector{UUID}=UUID[] + ignore_ids::Vector{UUID}=UUID[], + param_type::Type{<:AbstractParameterType}=JuMPParameterType ) where {T<:FileType} # Load full set df = load(input_file, T) @@ -191,12 +206,13 @@ function load(model_file::AbstractString, input_file::AbstractString, ::Type{T}; # No batch if isnothing(batch_size) pairs = _dataframe_to_dict(df, model_file) - return ProblemIterator(ids, pairs) + return ProblemIterator(pairs; ids=ids, param_type=param_type) end # Batch num_batches = ceil(Int, length(ids) / batch_size) idx_range = (i) -> (i-1)*batch_size+1:min(i*batch_size, length(ids)) - return (i) -> ProblemIterator(ids[idx_range(i)], _dataframe_to_dict(df[idx_range(i), :], model_file)), num_batches + return (i) -> ProblemIterator(_dataframe_to_dict(df[idx_range(i), :], model_file); + ids=ids[idx_range(i)], param_type=param_type), num_batches end """ @@ -204,18 +220,26 @@ end Update the value of a parameter in a JuMP model. """ -function update_model!(model::JuMP.Model, p::VariableRef, val) +function update_model!(::Type{POIParamaterType}, model::JuMP.Model, p::VariableRef, val) return MOI.set(model, POI.ParameterValue(), p, val) end +function update_model!(::Type{JuMPNLPParameterType}, model::JuMP.Model, p::VariableRef, val) + return set_parameter_value(p, val) +end + +function update_model!(::Type{JuMPParameterType}, model::JuMP.Model, p::VariableRef, val) + return fix(p, val) +end + """ update_model!(model::JuMP.Model, pairs::Dict, idx::Integer) Update the values of parameters in a JuMP model. """ -function update_model!(model::JuMP.Model, pairs::Dict, idx::Integer) +function update_model!(model::JuMP.Model, pairs::Dict, idx::Integer, param_type::Type{<:AbstractParameterType}) for (p, val) in pairs - update_model!(model, p, val[idx]) + update_model!(param_type, model, p, val[idx]) end end @@ -228,7 +252,8 @@ function solve_and_record( problem_iterator::ProblemIterator, recorder::Recorder, idx::Integer ) model = problem_iterator.model - update_model!(model, problem_iterator.pairs, idx) + problem_iterator.pre_solve_hook(model) + update_model!(model, problem_iterator.pairs, idx, problem_iterator.param_type) optimize!(model) status = recorder.filterfn(model) early_stop_bool = problem_iterator.early_stop(model, status, recorder) diff --git a/src/samplers.jl b/src/samplers.jl index e27db6d..a6649e5 100644 --- a/src/samplers.jl +++ b/src/samplers.jl @@ -119,7 +119,7 @@ end Load the parameters from a JuMP model. """ function load_parameters(model::JuMP.Model) - cons = constraint_object.(all_constraints(model, VariableRef, MOI.Parameter{Float64})) + cons = constraint_object.([all_constraints(model, VariableRef, MOI.Parameter{Float64}); all_constraints(model, VariableRef, MOI.EqualTo{Float64})]) parameters = [cons[i].func for i in 1:length(cons)] vals = [cons[i].set.value for i in 1:length(cons)] return parameters, vals diff --git a/test/datasetgen.jl b/test/datasetgen.jl index dee94de..8e284a8 100644 --- a/test/datasetgen.jl +++ b/test/datasetgen.jl @@ -5,7 +5,7 @@ Test dataset generation for different filetypes """ function test_problem_iterator(path::AbstractString) - @testset "Dataset Generation Type: $filetype" for filetype in [CSVFile, ArrowFile] + @testset "Dataset Generation (POI) Type: $filetype" for filetype in [CSVFile, ArrowFile] # The problem to iterate over model = JuMP.Model(() -> POI.Optimizer(HiGHS.Optimizer())) @variable(model, x) @@ -55,6 +55,17 @@ function test_problem_iterator(path::AbstractString) @test num_p * successfull_solves == 1 end + @testset "pre_solve_hook" begin + file_dual_output = joinpath(path, "test_$(string(uuid1()))_output") # file path + recorder_dual = Recorder{filetype}(file_dual_output; dual_variables=[cons]) + sum_p = 0 + problem_iterator = ProblemIterator( + Dict(p => collect(1.0:num_p)); pre_solve_hook=(args...) -> sum_p += 1 + ) + successfull_solves = solve_batch(problem_iterator, recorder_dual) + @test sum_p == num_p + end + @testset "solve_batch" begin successfull_solves = solve_batch(problem_iterator, recorder) @@ -96,6 +107,58 @@ function test_problem_iterator(path::AbstractString) end end end + @testset "Dataset Generation JuMP" begin + model = JuMP.Model(HiGHS.Optimizer) + @variable(model, x) + p = @variable(model, _p) + @constraint(model, cons, x + _p >= 3) + @objective(model, Min, 2x) + num_p = 10 + batch_id = string(uuid1()) + problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)); param_type=L2O.JuMPParameterType) + file_output = joinpath(path, "test_$(batch_id)_output") # file path + recorder = Recorder{ArrowFile}( + file_output; primal_variables=[x], dual_variables=[cons] + ) + successfull_solves = solve_batch(problem_iterator, recorder) + iter_files = readdir(joinpath(path)) + iter_files = filter(x -> occursin(string(ArrowFile), x), iter_files) + file_outs = [ + joinpath(path, file) for + file in iter_files if occursin("$(batch_id)_output", file) + ] + # test output file + df = Arrow.Table(file_outs) + @test length(df) == 8 + @test length(df[1]) == num_p * successfull_solves + rm.(file_outs) + end + @testset "Dataset Generation JuMPNLP" begin + model = JuMP.Model(Ipopt.Optimizer) + @variable(model, x) + p = @variable(model, _p in MOI.Parameter(1.0)) + @constraint(model, cons, x + _p >= 3) + @objective(model, Min, 2x) + num_p = 10 + batch_id = string(uuid1()) + problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)); param_type=L2O.JuMPNLPParameterType) + file_output = joinpath(path, "test_$(batch_id)_output") # file path + recorder = Recorder{ArrowFile}( + file_output; primal_variables=[x], dual_variables=[cons] + ) + successfull_solves = solve_batch(problem_iterator, recorder) + iter_files = readdir(joinpath(path)) + iter_files = filter(x -> occursin(string(ArrowFile), x), iter_files) + file_outs = [ + joinpath(path, file) for + file in iter_files if occursin("$(batch_id)_output", file) + ] + # test output file + df = Arrow.Table(file_outs) + @test length(df) == 8 + @test length(df[1]) == num_p * successfull_solves + rm.(file_outs) + end end function test_load(model_file::AbstractString, input_file::AbstractString, ::Type{T}, ids::Vector{UUID}; diff --git a/test/test_flux_forecaster.jl b/test/test_flux_forecaster.jl index 659585a..3a7be98 100644 --- a/test/test_flux_forecaster.jl +++ b/test/test_flux_forecaster.jl @@ -21,7 +21,7 @@ function test_flux_forecaster(file_in::AbstractString, file_out::AbstractString) rng=123, epochs=20, optimiser=ConvexRule( - Flux.Optimise.Adam(0.001, (0.9, 0.999), 1.0e-8, IdDict{Any,Any}()) + Optimisers.Adam() ), )