Custom layer for a good guess of system dynamics

General Usage
1

Hi, I’m trying to create a NeuralODE for the forcasting of a dynamic system.
I did a NN for the forcasting of the system under different initial conditions and it’s work very well.

Now I would like to create NN capable to learn the infuences of a possible variation of a input value, to do that I’m trying to incorporate in the first layer of the NN a customized layer that take cares of the dynamic of the system I want to reproduce.

The problem is that I don’t understand how can I bring the parameters of this custom layer together with the other parameters of the whole NN.

function law(u,p,t)
    
    F = Float32(5.0)
    α = p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return[z1,z2]
end

#nn
NN = Lux.Chain(x -> law(x,p,t),
                Lux.Dense(2,20,tanh),
                Lux.Dense(30,10,tanh),
                Lux.Dense(10,2))
p, st = Lux.setup(rng, NN)

prob_neuralode = NeuralODE(NN, tspan, Tsit5(), saveat = t)
p_init = Lux.ComponentArray(p)

the output is always an empty first layer

ComponentVector{Float32}(layer_1 = Float32[], layer_2 = (weight = Float32[0.567042 0.4250619; 0.042177025 0.5778097; … ; -0.30559936 -0.09352723; 0.20134257 -0.40837312], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_3 = (weight = Float32[-0.1402396 0.52604 … -0.17828122 0.22984686; -0.20939533 0.24967496 … 0.2779108 -0.4940759; … ; -0.44822726 0.3582243 … -0.25572324 -0.05863302; 0.43971282 0.09462769 … -0.12614584 -0.41752148], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_4 = (weight = Float32[0.36662766 0.41478822 … -0.37065506 -0.43711254; -0.15670164 -0.54979855 … -0.58355767 0.009892202], bias = Float32[0.0; 0.0]))
2

The problem here is that Lux.setup doesn’t know how to initialize p for your custom layer so it assumes that your layer has no parameters. See Lux Interface - Lux.jl for how to make custom layers work.

2 Likes
3

What’s the error message thrown on that? This seems something that could be handled by posting exactly that when such a case is seen?

1 Like
4

The issue with the definition OP posted is that it is a perfectly valid model declaration (Chain sees that it is a function with 1 argument, so it is automatically converted to a layer with no parameter/state internally), So an error can only occur at runtime when it sees that p and t are undefined.

1 Like
5

There is no error in this problem formulation. I can’t figure out how can I pass the parameters of the approximated dynamic layer law(u,p,t) in the NeuralODE and soo on in the optimization problem.

Any idea?

6 
struct LawLayer <: Lux.AbstractExplicitLayer
end

Lux.initialparameters(rng, l::LawLayer) = (p = define how to initialize p,)

function (::LawLayer)(u, ps, st)
    
    F = Float32(5.0)
    α = ps.p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return ([z1,z2], st)
end

See Training a Simple LSTM - Lux.jl

7

Really thank you for your support now I’ve clarified several doubts but I miss something yet, please be patient.
I tried to follow this tutorial Lux Interface.

struct LawLayer{F1} <: Lux.AbstractExplicitLayer
  parameters::F1
end

function LawLayer(; parameters  = Lux.glorot_uniform)
return LawLayer{typeof(parameters)}(parameters)
end

LawLayer()

Lux.initialparameters(rng, l::LawLayer) = (p = l.parameters(rng,1,1))
Lux.initialstates(::AbstractRNG, ::LawLayer) = NamedTuple()
Lux.parameterlength(::LawLayer) = 1
Lux.statelength(::LawLayer) = 0

function (::LawLayer)(u, ps, st)
    
    F = Float32(5.0)
    α = ps.p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return ([z1,z2], st)
end

but I recive the same result:

Chain(
    layer_1 = WrappedFunction(#8),
    layer_2 = Dense(2 => 30, tanh_fast),  # 90 parameters
    layer_3 = Dense(30 => 30, tanh_fast),  # 930 parameters
    layer_4 = Dense(30 => 2),           # 62 parameters
)         # Total: 1_082 parameters,
          #        plus 0 states, summarysize 48 bytes.

((layer_1 = NamedTuple(), layer_2 = (weight = Float32[0.10518762 0.30059525; -0.27172458 -0.36132228; … ; 0.061005663 0.044980295; -0.2633441 0.023604192], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_3 = (weight = Float32[0.16688053 -0.042258434 … 0.07908167 -0.19162428; 0.19909461 0.16200048 … 0.11714643 -0.03535033; … ; -0.13860366 0.009593734 … -0.066989586 -0.0815194; 0.21782668 -0.23353975 … -0.15897466 0.22407342], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_4 = (weight = Float32[0.215354 0.35274163 … 0.15960616 -0.2909302; 0.055675678 0.3207232 … -0.34296218 -0.34948933], bias = Float32[0.0; 0.0])), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple()))

what I’m wrong?

8

How are you defining the Chain? Seems like layer_1 is a function and not LawLayer()

9

Sorry, here the whole code:

rng = Random.default_rng()

in_layer = 30
hidden_layer = 30 

struct LawLayer{F1} <: Lux.AbstractExplicitLayer
  parameters::F1
end

function LawLayer(; parameters  = Lux.glorot_uniform)
return LawLayer{typeof(parameters)}(parameters)
end

LawLayer()

Lux.initialparameters(rng, l::LawLayer) = (p = l.parameters(rng,1,1))
Lux.initialstates(::AbstractRNG, ::LawLayer) = NamedTuple()
Lux.parameterlength(::LawLayer) = 1
Lux.statelength(::LawLayer) = 0

function (::LawLayer)(u, ps, st)
    
    F = Float32(5.0)
    α = ps.p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return ([z1,z2], st)
end

#nn
NN = Lux.Chain(u -> LawLayer(u,ps,st),
                Lux.Dense(2,in_layer,tanh),
                Lux.Dense(in_layer,hidden_layer,tanh),
                Lux.Dense(hidden_layer,2))
                      
p, st = Lux.setup(rng, NN)

prob_neuralode = NeuralODE(NN, tspan, Tsit5(), saveat = t)
p_init = Lux.ComponentArray(p)
10

Change u -> LawLayer(u,ps,st) to LawLayer()

1 Like