# Nonlinear fit with Flux

**URL:** <https://discourse.julialang.org/t/nonlinear-fit-with-flux/53075>\
**Category:** Machine Learning\
**Tags:** flux\
**Created:** [January 9, 2021, 1:12am UTC](https://discourse.julialang.org/t/nonlinear-fit-with-flux/53075 "2021-01-09T01:12:14Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![edwinb-ai](https://avatars.discourse-cdn.com/v4/letter/e/bb73d2/32.png) [@edwinb-ai](https://discourse.julialang.org/u/edwinb-ai)\
**Post date:** [January 9, 2021, 1:12am UTC](https://discourse.julialang.org/t/nonlinear-fit-with-flux/53075/1 "2021-01-09T01:12:14Z")

</div>

```julia
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](https://global.discourse-cdn.com/julialang/original/3X/5/d/5d8dd7d0fc877b148b9a94e96e95aa0a78747e1c.png)

Any advice would be greatly appreciated.

---

<div class="post-metadata">

**Author:** ![jstrube](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jstrube/32/525_2.png) [@jstrube](https://discourse.julialang.org/u/jstrube)\
**Post date:** [January 9, 2021, 3:29am UTC](https://discourse.julialang.org/t/nonlinear-fit-with-flux/53075/2 "2021-01-09T03:29:58Z")

</div>

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

```julia
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.

---

<div class="post-metadata">

**Author:** ![edwinb-ai](https://avatars.discourse-cdn.com/v4/letter/e/bb73d2/32.png) [@edwinb-ai](https://discourse.julialang.org/u/edwinb-ai)\
**Post date:** [January 10, 2021, 4:10am UTC](https://discourse.julialang.org/t/nonlinear-fit-with-flux/53075/3 "2021-01-10T04:10:43Z")

</div>

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.