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
Any advice would be greatly appreciated.