fix Flux code to version 0.16.3 and update deps

Author: Giovanni3A

Project.toml

@@ -11,18 +11,20 @@ JobQueueMPI = "32d208e1-246e-420c-b6ff-18b71b410923"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 BilevelJuMP = "0.6.2"
 DiffOpt = "0.5.0"
-Flux = "0.14.25"
+Flux = "0.16.3"
 JobQueueMPI = "0.1.1"
 JuMP = "1.24"
-Optim = "1.11"
 MPI = "0.20.22"
+Optim = "1.11"
+Optimisers = "0.4.5"
 ParametricOptInterface = "0.9.0"
 Zygote = "0.6.75"
 julia = "1.10"

src/optimizers/gradient.jl

@@ -42,6 +42,7 @@ function train_with_gradient!(
     best_θ = extract_params(model.forecast)
     trace = Array{Float64}(undef, epochs)
     stochastic = batch_size > 0
+    opt_state = Flux.setup(rule, model.forecast)
 
     # precompute batches
     batches = repeat(1:T, outer = (1, epochs))'
@@ -87,7 +88,7 @@ function train_with_gradient!(
         end
 
         # take gradient step
-        apply_gradient!(model.forecast, dC, epochx, rule)
+        apply_gradient!(model.forecast, dC, epochx, opt_state)
     end
 
     # fix best model

src/optimizers/gradient_mpi.jl

@@ -32,6 +32,7 @@ function train_with_gradient_mpi!(
     T = size(X)[1]
     stochastic = batch_size > 0
     compute_full_cost = true
+    opt_state = Flux.setup(rule, model.forecast)
 
     # precompute batches
     batches = repeat(1:T, outer = (1, epochs))'
@@ -121,7 +122,7 @@ function train_with_gradient_mpi!(
             end
 
             # take gradient step (if not last epoch)
-            apply_gradient!(model.forecast, dCdy, epochx, rule)
+            apply_gradient!(model.forecast, dCdy, epochx, opt_state)
         end
 
         # release workers

src/predictive_model.jl

@@ -1,6 +1,7 @@
 using Flux
 using Statistics
 import Zygote
+import Optimisers
 
 """
     PredictiveModel(networks, input_output_map, input_size, output_size)
@@ -198,7 +199,7 @@ function apply_params(model::PredictiveModel, θ)
 end
 
 """
-    apply_gradient!(model, dCdy, X, optimizer)
+    apply_gradient!(model, dCdy, X, rule)
 
 Apply a gradient vector to the model parameters.
 
@@ -209,17 +210,16 @@ Apply a gradient vector to the model parameters.
   - `model::PredictiveModel`: model to be updated.
   - `dCdy::Vector{<:Real}`: gradient vector.
   - `X::Matrix{<:Real}`: input data.
-  - `optimizer`: Optimiser to be used.
+  - `rule`: Optimisation rule.
     ...
 """
 function apply_gradient!(
     model::PredictiveModel,
     dCdy::Vector{<:Real},
     X::Matrix{<:Real},
-    optimizer,
+    opt_state,
 )
-    ps = Flux.params(model.networks)
-    loss(x, y) = mean(dCdy'model(x))
-    train_data = [(X', 0.0)]
-    return Flux.train!(loss, ps, train_data, optimizer)
+    loss3(m, X) = mean(dCdy'm(X'))
+    grad = Zygote.gradient(loss3, model, X)[1]
+    Optimisers.update!(opt_state, model, grad)
 end

test/Project.toml

@@ -8,6 +8,7 @@ JobQueueMPI = "32d208e1-246e-420c-b6ff-18b71b410923"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

test/test_predictive_model.jl

@@ -24,7 +24,7 @@ out_size = 2
         forecaster,
         ones(out_size),
         ones((1, in_size)),
-        Flux.Descent(0.1),
+        Flux.setup(Flux.Descent(0.1), forecaster),
     )
     @test Flux.params(forecaster.networks[1])[1] ==
           0.9 * ones((out_size, in_size))
@@ -53,7 +53,7 @@ end
         forecaster,
         ones(out_size),
         ones((1, in_size)),
-        Flux.Descent(0.1),
+        Flux.setup(Flux.Descent(0.1), forecaster),
     )
     @test Flux.params(forecaster.networks[1])[1] ==
           0.9 * ones((out_size, in_size))
@@ -84,7 +84,7 @@ end
         forecaster,
         ones(out_size),
         ones((1, in_size)),
-        Flux.Descent(0.1),
+        Flux.setup(Flux.Descent(0.1), forecaster),
     )
     @test Flux.params(forecaster.networks[1])[1] ==
           0.8 * ones((model_out_size, model_in_size))
@@ -122,7 +122,7 @@ end
         forecaster,
         ones(out_size),
         ones((1, in_size)),
-        Flux.Descent(0.1),
+        Flux.setup(Flux.Descent(0.1), forecaster),
     )
     @test Flux.params(forecaster.networks[1])[1] ==
           0.9 * ones((model_out_size, model_in_size))