demo of lm optimization

lm.jl

using LsqFit
using ForwardDiff
using Plots
using Flux
using MINPACK

xdata = 0.0:0.1:10.0
xdata = xdata'
ydata = Float64.(vec(sin.(xdata)))

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

function f!(fvec, p)
    ypred = vec(re(p)(xdata))
    @. fvec = abs(ypred - ydata)
    return fvec
end
fvec = similar(ydata)
f!(fvec, p)

function g!(fjac, p)
    ForwardDiff.jacobian!(fjac, (fvec, p) -> f!(fvec, p), fvec, p)
    return fjac
end
fjac = zeros(length(ydata), length(p))
g!(fjac, p)

p0 = Float64.(p);
m = length(ydata)
res = fsolve(
    f!,
    g!,
    p0,
    m,
    iterations = 300,
    tol = 1e-8,
    show_trace = true,
    tracing = true;
    method = :lm,
)

l_loss = ones(res.trace.f_calls, 4)
for i = 1:res.trace.f_calls
    l_loss[i, 1] = res.trace.trace[i].iteration
    l_loss[i, 2] = res.trace.trace[i].fnorm
    l_loss[i, 3] = res.trace.trace[i].xnorm
    l_loss[i, 4] = res.trace.trace[i].step_time
end

l_plt = []
plt = scatter(xdata[:], ydata[:], label = "data");
plot!(plt, xdata[:], vec(re(res.x)(xdata)), label = "pred");
push!(l_plt, plt)

plt = plot(
    l_loss[:, 1],
    l_loss[:, 2],
    xscale = :identity,
    yscale = :log10,
    label = "fnorm",
);
push!(l_plt, plt)
plt = plot(
    l_loss[:, 1],
    l_loss[:, 3] .+ 1.e-6,
    xscale = :identity,
    yscale = :log10,
    label = "xnorm",
);
push!(l_plt, plt)

plt = plot(l_plt...)
png(plt, "lmnn")

Using LsqFit.jl

using LsqFit
using ForwardDiff
using Plots
using Flux

xdata = 0.0:0.1:10.0
xdata = xdata'
ydata = Float64.(vec(sin.(xdata)))

nn = Chain(
        Dense(1, 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)

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

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

display(res.param)

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