How do I create arbitrary parameterized layers in DiffEqFlux.lj neuralODE? Julia Julialang Flux.jl - neural-network

Question

I am able to create and optimize neuralODEs in julia(1.3 and 1.2) using Flux.jl and DiffEqFlux.jl but it fails under a crucial important general case.
what works:
I can train the Neural net parameters if it is built out of the
provided Flux.jl layers like Dense().
I can include an arbitrary function as a layer in the network chain, e.g. x -> x.*x
What fails:
However if the arbitrary function has parameters I want to train then Flux. Train will not adjust these parameters causing it to fail.
I have tried making these added parameters Tracked and included in the list of parameters given to the training system but it ignores them and they remain unvaried.
The documentation says very cryptically that one can use Flux.#functor on a layer to make sure it's parameters get tracked. However functor was not included in Flux till version 0.10.0 and the only version of Flux compatible with NeuralODEs in DiffEqFlux is 0.9.0
So here's an toy example of a 2 layer neural net I want to use
p = param([1.0])
dudt = chain( x -> p[1]*x.*x, Dense(2,2) )
ps = Flux.params(dudt)
then I use the flux train on this. when I do this the parameter p is not varied, but the parameters in the Dense layer are.
I have tried explicitly including like this
ps = Flux.Params([p,dudt])
but that has the same result and the same problem
I think what I need to do is build a struct with an associted function that implements the
x->p[1]*x*x
then call #functor on this. That struct can then be used in the chain.
But as I noted the version of Flux with #functor is not compatible with DiffEqFlux of any version.
So I need a way to make flux pay attention to my custom parameters, not just the ones in Dense()
How???

Answer 1

I think I get what your question is, but please clarify if I am answering the wrong question here. The issue is that the p is only grabbing from a global reference and thus not differentiated during adjoints. A much better way to handle this in 2020 is to use FastChain. The FastChan interface lets you define layer functions and their parameter dependencies, so this is a nice way to make your neural network incorporate arbitrary functions with parameters. Here's what that looks like:
using DifferentialEquations
using Flux, Zygote
using DiffEqFlux
x = Float32[2.; 0.]
p = Float32[2.0]
tspan = (0.0f0,1.0f0)
mylayer(x,p) = p[1]*x
DiffEqFlux.paramlength(::typeof(mylayer)) = 1
DiffEqFlux.initial_params(::typeof(mylayer)) = rand(Float32,1)
dudt = FastChain(FastDense(2,50,tanh),FastDense(50,2),mylayer)
p = DiffEqFlux.initial_params(dudt)
function f(u,p,t)
dudt(u,p)
end
ex_neural_ode(x,p) = solve(ODEProblem(f,x,tspan,p),Tsit5())
solve(ODEProblem(f,x,tspan,p),Tsit5())
du0,dp = Zygote.gradient((x,p)->sum(ex_neural_ode(x,p)),x,p)
where the last value of p is the one parameter for p in mylayer. Or you can directly use Flux:
using DifferentialEquations
using Flux, Zygote
using DiffEqFlux
x = Float32[2.; 0.]
p2 = Float32[2.0]
tspan = (0.0f0,1.0f0)
dudt = Chain(Dense(2,50,tanh),Dense(50,2))
p,re = Flux.destructure(dudt)
function f(u,p,t)
re(p[1:end-1])(u) |> x-> p[end]*x
end
ex_neural_ode() = solve(ODEProblem(f,x,tspan,[p;p2]),Tsit5())
grads = Zygote.gradient(()->sum(ex_neural_ode()),Flux.params(x,p,p2))
grads[x]
grads[p]
grads[p2]