Merge pull request #28 from LAMPSPUC/dev

Dev

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/flux_utils.jl

@@ -7,7 +7,7 @@ Extract the parameters of a Flux model (Flux.Chain or Flux.Dense) into a single
 vector.
 """
 function extract_flux_params(model::Union{Flux.Chain,Flux.Dense})
-    θ = Flux.params(model)
+    θ = Flux.trainables(model)
     return reduce(vcat, [vec(p) for p in θ])
 end
 
@@ -21,7 +21,7 @@ function fix_flux_params_single_model(
     θ::Vector{<:Real},
 )
     i = 1
-    for p in Flux.params(model)
+    for p in Flux.trainables(model)
         psize = prod(size(p))
         p .= reshape(θ[i:i+psize-1], size(p))
         i += psize
@@ -38,7 +38,7 @@ of parameters.
 function fix_flux_params_multi_model(models, θ::Vector{<:Real})
     i = 1
     for model in models
-        for p in Flux.params(model)
+        for p in Flux.trainables(model)
             psize = prod(size(p))
             p .= reshape(θ[i:i+psize-1], size(p))
             i += psize
@@ -54,8 +54,8 @@ Check if a Flux layer has parameters.
 """
 function has_params(layer)
     try
-        # Attempt to get parameters; if it works and isn't empty, return true
-        return !isempty(Flux.params(layer))
+        # Attempt to get trainable parameters; if it works and isn't empty, return true
+        return !isempty(Flux.trainable(layer))
     catch e
         # If there is an error (e.g. method not matching), assume no parameters
         return false

src/optimizers/bilevel.jl

@@ -177,10 +177,10 @@ function solve_bilevel(
     ilayer = 1
     for layer in model.forecast.networks[1]
         if has_params(layer)
-            for p in Flux.params(layer.weight)
+            for p in Flux.trainables(layer.weight)
                 p .= value.(predictive_model_vars[ilayer][:W])
             end
-            for p in Flux.params(layer.bias)
+            for p in Flux.trainables(layer.bias)
                 p .= value.(predictive_model_vars[ilayer][:b])
             end
         end

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]
+    return 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,11 +24,10 @@ 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))
-    @test Flux.params(forecaster.networks[1])[2] == 0.9 * ones(out_size)
+    @test Flux.trainables(forecaster)[1] == 0.9 * ones((out_size, in_size))
+    @test Flux.trainables(forecaster)[2] == 0.9 * ones(out_size)
 end
 
 @testset "Single-Chain" begin
@@ -53,11 +52,10 @@ 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))
-    @test Flux.params(forecaster.networks[1])[2] == 0.9 * ones(out_size)
+    @test Flux.trainables(forecaster)[1] == 0.9 * ones((out_size, in_size))
+    @test Flux.trainables(forecaster)[2] == 0.9 * ones(out_size)
 end
 
 @testset "Multi-Variate-Dense" begin
@@ -84,11 +82,11 @@ 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] ==
+    @test Flux.trainables(forecaster)[1] ==
           0.8 * ones((model_out_size, model_in_size))
-    @test Flux.params(forecaster.networks[1])[2] == 0.8 * ones(model_out_size)
+    @test Flux.trainables(forecaster)[2] == 0.8 * ones(model_out_size)
 end
 
 @testset "Multi-Model-Dense" begin
@@ -122,12 +120,12 @@ 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] ==
+    @test Flux.trainables(forecaster)[1] ==
           0.9 * ones((model_out_size, model_in_size))
-    @test Flux.params(forecaster.networks[1])[2] == 0.9 * ones(model_out_size)
-    @test Flux.params(forecaster.networks[2])[1] ==
+    @test Flux.trainables(forecaster)[2] == 0.9 * ones(model_out_size)
+    @test Flux.trainables(forecaster)[3] ==
           0.9 * ones((model_out_size, model_in_size))
-    @test Flux.params(forecaster.networks[2])[2] == 0.9 * ones(model_out_size)
+    @test Flux.trainables(forecaster)[4] == 0.9 * ones(model_out_size)
 end