Nonlinear fit with Flux

using Flux
using Flux.Losses: mse
using Flux.Data: DataLoader
using Plots
using Random


nonlinear(x) = sin(x) + sin(2 * x) + 0.1f0 * randn()

function get_data()
    points = 1000
    xvals = collect(Float32, range(-10.0, 10.0, length=points))
    yvals = nonlinear.(xvals)
    xvals = reshape(xvals, 1, :)

    return xvals, yvals

function build_model()
    hidden = 64

    return Flux.Chain(
        Flux.Dense(1, hidden),
        Flux.Dense(hidden, hidden, Flux.elu),
        Flux.Dense(hidden, 1)

function train(xvals, yvals)
    model = build_model()

    loss(x, y) = mse(model(x)', y)

    trainer = DataLoader((xvals, yvals), shuffle=true, batchsize=32)

    @Flux.epochs 300 Flux.train!(loss, params(model), trainer, Flux.Descent(0.001))

    return model

model = train(xvals, yvals)

scatter(xvals', yvals, ms=2.5, label="Data")
display(mse(model(xvals)', yvals))
plot!(xvals', model(xvals)', linewidth=2, label="Fit")

I want to use the previous code to fit a simple nonlinear function with a neural network, in order to practice with Flux.jl, but I’m getting terrible perfomance.
This is the result I obtain


Any advice would be greatly appreciated.

I don’t see a problem. The model might be a bit too shallow.
If I change the model to

        Flux.Dense(1, hidden, Flux.elu),
        Flux.Dense(hidden, hidden, Flux.elu),
        Flux.Dense(hidden, hidden, Flux.elu),
		Flux.Dense(hidden, 1)

I get a nice fit after 3000 epochs. Not sure how shallow / how few epochs are the minimum.

1 Like

Thank you!

This advice worked wonders.
And it also made me realize some other stuff is wrong, such as sending the whole training procedure to the GPU. But that is another issue.