Dear Forum,
I am trying to add ODE into machine learning pipeline and calculate the gradient of an ode layer that takes the RHS as an argument. We tried to feed the RHS, in terms of inputFunction(t), into ODE (i.e. fiip function). In fiip, inputFunction(t) is extracted from p, which is fed in function model(inputdata).
The code is as below.
using Flux
using DifferentialEquations, DiffEqSensitivity
parameter = [0.5]
function fiip(du, u, p, t)
p1, inputFunction = p
du[1] = dx = p1 * u[1] + inputFunction(t)
du[2] = dy = -p1 * u[2]
end
function model(inputdata)
inputFunction(t) = inputdata * t
p = [parameter[1], inputFunction]
u0 = [1.0; 1.0];
tspan = (0.0, 10.0)
concrete_solve(
ODEProblem(fiip, u0, tspan, p),
Tsit5(),
u0,
p,
saveat = 0:0.1:10,
sensealg = QuadratureAdjoint(),
)[1,:]
end
loss(x, y) = Flux.mse(model(x), y) # Loss function
x = 10
y = 100 * sin.(0:0.1:10)
# Gradient calculation
gs = Flux.gradient(() -> loss(x, y), params(parameter))
When the code executes to the last line, i.e. calculate the gradient of ODE with respect to the designated parameter, such error pops up:
LoadError: TypeError: in TrackedReal, in V, expected V<:Real, got Type{Any}
We’d be highly grateful for any insight into the possible origin and solution to this problem.
Also, will it help if we use callbacks on the ODE solver instead?