fit nn using LsqFit.jl

lm_nn.jl

using LsqFit
using ForwardDiff
using Plots
using Flux

xdata = 0.0:1.0:10.0
xdata = xdata'
ydata = Float32.(vec(sin.(xdata)))

nn = Chain(
        Dense(1, 5, tanh),
        Dense(5, 5, tanh),
        Dense(5, 1))
x, re = Flux.destructure(nn)

function f(p)
    ypred = vec(re(p)(xdata))
    return @. abs(ypred - ydata)
end
f(x)

g = function (p)
    return ForwardDiff.jacobian(x -> f(x), p)
end
g(x)

fit = LsqFit.lmfit(f, g, x, Float64[]; show_trace=true, maxIter=1000, x_tol=1e-6)

display(fit.param)

plt = scatter(xdata[:], ydata[:], label="data");
plot!(plt, xdata[:], vec(re(fit.param)(xdata)), label="pred");
png(plt, "lmnn")