Adding external parameters to neural network

Specific Domains Machine Learning
, ,
2

yes you can. Set it up like https://lux.csail.mit.edu/stable/tutorials/intermediate/1_NeuralODE#Define-the-Neural-ODE-Layer. Here is a pseudocode

# `c` will have to be an array for it to be trainable
my_custom_model = @compact(; model, c = [2.0]) do x, ps
   # Note that state handling is automatic with `@compact`
    function rhs!(du, u, p, t)
       û = model(u, p.model)  # Parameters accessible via p.<fieldname> 
       du[1] = û[1] + log.(p.c[1]./u[1])  # p.c[1] since we set it up as an array
       du[2] = û[2] 
       du[3] = û[3]
    end

    prob = ....
    return solve(prob, ....)
end

ps, st = Lux.setup(rng, my_custom_model)
# ps.c and ps.model are automatically populated

my_custom_model(x_input, ps, st)

There is a new tutorial that is going to be merged this week, which demonstrates how to do UDEs with a clean syntax Solving Optimal Control Problems with Symbolic Universal Differential Equations | Lux.jl Documentation

2 Likes