Unable to evaluate a Lux model at a specific parameter set.

[TorkelE]

Say that I have a Lux model

using Lux
model = Lux.Chain(
    Lux.Dense(1 => 4, Lux.softplus),
    Lux.Dense(4 => 4, Lux.softplus),
    Lux.Dense(4 => 1, Lux.softplus)
)

Two questions:

1. is there a way to learn how many parameters this model have? All function I apply returns various structured things which I cannot really understand what they contain.
2. Lets say my model has 5 parameters and 1 input, how do I, given that I have these in normal vectors:

input = [1.0]
p_vals = [1.0, -2.0, 0.5, 5.0, -1.5]

get the output of my neural network for these parameter values and output?

I have spent the last couple of days trying to get something like this to work, but I really have not been able to. i have been trying to make a basic example where I can train a model using optimization (where I have a neural network and some other things I want to optimize), and then build from there. Unfortunatley I cannot get this to work. Doc pages I have been reading include
https://lux.csail.mit.edu/stable/introduction/
https://lux.csail.mit.edu/stable/tutorials/beginner/1_Basics
https://docs.sciml.ai/Optimization/stable/optimization_packages/optimization/#Train-NN-with-Sophia
https://lux.csail.mit.edu/stable/tutorials/beginner/5_OptimizationIntegration
https://lux.csail.mit.edu/stable/tutorials/beginner/2_PolynomialFitting

[TorkelE]

(I realise that I might need to create a component array from my vector, if that is the case I don't mind if that is required)

[avik-pal]

For the interface check https://lux.csail.mit.edu/stable/manual/interface#Layer-Interface.

1. Use parameterlength for the number of parameters.
2. You should construct an initial random parameter set using Lux.setup and update those parameter values with the values you want. https://lux.csail.mit.edu/stable/manual/interface#Parameter-Interface describes the expectations from the parameters struct.

[TorkelE]

Thanks a lot, I think I am almost there now. Using

model = Lux.Chain(
    Lux.Dense(1 => 4, Lux.softplus),
    Lux.Dense(4 => 4, Lux.softplus),
    Lux.Dense(4 => 1, Lux.softplus)
)
rng = Random.default_rng()
ps, st = Lux.setup(rng, model)

I get a good starting point. Cheking, my parameters are in a layered named tuples where e.g.

ps.layer_1.weight

is a 4-valued vector (I presume with all the weights of the 4 layer).

Is there a direct interface to construct a correct such a layered NamedTuple from a vector with parameters values? I might be able to get there my self, but think there is a decent chance I'd make a mistake, and there is a quite a bit of hard coding to do.

The reason that I want to work with normal vectors is that that is what Optimization uses, so if I have an OptimizationProblem I want to be able to convert the values of that one to something that Lux can use. I think there are some cases where Optimization can work with Lux directly, but I could not get that to work for my purpose (where some of the parameters are not related to neural networks, i.e. a bit like an UDE, but I am working on a different case)

[TorkelE]

Ok, I think I have figured out how to do it, thanks a lot for the help.

Optimization can work over ComponentArrays, and Lux ignores fields in the input that are not related to the neural network. So I can simply create a component array for my neural network parameters, and then expand it with whatever other stuff I want my optimization to consider. Then I can input this into a loss function and then a OptimizationProblem.

# Fetch packages
using Optimization,Lux, Zygote, Plots, Random, ComponentArrays
import OptimizationOptimisers: Adam

# Create training data and true function (which we want the neural network to approximate).
f(x) = 0.5 * (x^3)/(x^3 + 0.7^3)
xs = collect(0.0:0.1:4.0)
ys = [(0.9 + 0.2rand())*f(x) for x in xs]

# Create the model.
model = Lux.Chain(
    Lux.Dense(1 => 4, Lux.softplus),
    Lux.Dense(4 => 4, Lux.softplus),
    Lux.Dense(4 => 1, Lux.softplus)
)
ps, st = Lux.setup(Random.default_rng(), model)

# Check input/output.
ps_ca = ComponentArray(ps)
input = reshape(xs, 1, length(xs))
target = reshape(ys, 1, length(xs))
ypred = LuxCore.stateless_apply(model, reshape(xs, 1, length(xs)), ps)

# Adds a non-neural network related field we want to optimise in parallel.
ps_ca_expanded = ComponentArray(ps_ca; new = 10.0)

# Defines loss function amd run optimization.
function loss(ps, (input, target))
    ypred = LuxCore.stateless_apply(model, input, ps)
    error = sum(abs2, ypred .- target)
    return error + abs(ps.new)
end
optf = OptimizationFunction(loss, AutoZygote())
prob = OptimizationProblem(optf, ps_ca_expanded, (input, target))
@time res = Optimization.solve(prob, Adam(0.001); maxiters = 100000)

# Checks output.
plot(f, xs[1], xs[end])
plot!(xs, ys; seriestype = :scatter)
ypred = LuxCore.stateless_apply(model, input, res.u)
plot!(xs, ypred[1,:]; label = "Predicted", color = :red)
res.u.new # should be close to `0.0`.