Nonlinear fit with Flux

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

gr()

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
end

function build_model()
    hidden = 64

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

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
end

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")
savefig("fit.png")

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

fit

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.Chain(
        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.