I wish to use Flux.training to optimize a parameter in parallel ensemble simulation, but it shows adjoint error as below:
ERROR: LoadError: Need an adjoint for constructor EnsembleSolution{Float64,2,Array{ODESolution{Float64,1,Array{Float64,1},Nothing,Nothing,Array{Float64,1},Array{Array{Float64,1},1},ODEProblem{Float64,Tuple{Float64,Float64},false,DiffEqBase.NullParameters,ODEFunction{false,var"#189#190",UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,var"#189#190",UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},OrdinaryDiffEq.Tsit5ConstantCache{Float64,Float64}},DiffEqBase.DEStats},1}}. Gradient is of type Array{Float64,2}
Does it mean Flux training is not compatible with EnsembleProblem? I’d be grateful of any insight into this problem solution.
The code is as below:
using DifferentialEquations, Flux
pa = [1.0]
function model1(input)
prob = ODEProblem((u, p, t) -> 1.01u * pa[1], 0.5, (0.0, 1.0))
function prob_func(prob, i, repeat)
remake(prob, u0 = rand() * prob.u0)
end
ensemble_prob = EnsembleProblem(prob, prob_func = prob_func)
sim = solve(ensemble_prob, Tsit5(), EnsembleThreads(), trajectories = 100)
end
Input_time_series = zeros(5, 100)
# loss function
loss(x, y) = Flux.mse(model1(x), y)
data = Iterators.repeated((Input_time_series, 0), 1)
cb = function () # callback function to observe training
println("Tracked Parameters: ", params(pa))
end
opt = ADAM(0.1)
println("Starting to train")
Flux.@epochs 10 Flux.train!(loss, params(pa), data, opt; cb = cb)