wehub-resource-sync
303 lines
9.4 KiB
Julia
303 lines
9.4 KiB
Julia
# Optimization in Julia. GradientDescent, SGD+Momentum, and Adam
|
|
# implemented as mutable structs with a common `step!` method.
|
|
# Driven on the Rosenbrock and saddle-point functions to show
|
|
# convergence, divergence, and saddle escape behavior.
|
|
# Stdlib only. Sources:
|
|
# https://docs.julialang.org/en/v1/manual/types/#Composite-Types
|
|
# https://arxiv.org/abs/1412.6980 (Adam: Kingma & Ba)
|
|
|
|
using Printf
|
|
|
|
|
|
abstract type Optimizer end
|
|
|
|
|
|
mutable struct GradientDescent <: Optimizer
|
|
lr::Float64
|
|
end
|
|
GradientDescent(; lr::Float64=0.001) = GradientDescent(lr)
|
|
|
|
function step!(opt::GradientDescent, params::Vector{Float64}, grads::Vector{Float64})
|
|
return params .- opt.lr .* grads
|
|
end
|
|
|
|
|
|
mutable struct SGDMomentum <: Optimizer
|
|
lr::Float64
|
|
momentum::Float64
|
|
velocity::Vector{Float64}
|
|
end
|
|
SGDMomentum(; lr::Float64=0.001, momentum::Float64=0.9) =
|
|
SGDMomentum(lr, momentum, Float64[])
|
|
|
|
function step!(opt::SGDMomentum, params::Vector{Float64}, grads::Vector{Float64})
|
|
if isempty(opt.velocity)
|
|
opt.velocity = zeros(length(params))
|
|
end
|
|
opt.velocity .= opt.momentum .* opt.velocity .+ grads
|
|
return params .- opt.lr .* opt.velocity
|
|
end
|
|
|
|
|
|
mutable struct Adam <: Optimizer
|
|
lr::Float64
|
|
beta1::Float64
|
|
beta2::Float64
|
|
epsilon::Float64
|
|
m::Vector{Float64}
|
|
v::Vector{Float64}
|
|
t::Int
|
|
end
|
|
Adam(; lr::Float64=0.001, beta1::Float64=0.9, beta2::Float64=0.999,
|
|
epsilon::Float64=1e-8) =
|
|
Adam(lr, beta1, beta2, epsilon, Float64[], Float64[], 0)
|
|
|
|
function step!(opt::Adam, params::Vector{Float64}, grads::Vector{Float64})
|
|
if isempty(opt.m)
|
|
opt.m = zeros(length(params))
|
|
opt.v = zeros(length(params))
|
|
end
|
|
opt.t += 1
|
|
opt.m .= opt.beta1 .* opt.m .+ (1 - opt.beta1) .* grads
|
|
opt.v .= opt.beta2 .* opt.v .+ (1 - opt.beta2) .* grads .^ 2
|
|
m_hat = opt.m ./ (1 - opt.beta1 ^ opt.t)
|
|
v_hat = opt.v ./ (1 - opt.beta2 ^ opt.t)
|
|
return params .- opt.lr .* m_hat ./ (sqrt.(v_hat) .+ opt.epsilon)
|
|
end
|
|
|
|
|
|
rosenbrock(p::Vector{Float64})::Float64 = (1 - p[1]) ^ 2 + 100 * (p[2] - p[1] ^ 2) ^ 2
|
|
|
|
|
|
function rosenbrock_grad(p::Vector{Float64})::Vector{Float64}
|
|
x, y = p[1], p[2]
|
|
df_dx = -2 * (1 - x) + 200 * (y - x ^ 2) * (-2 * x)
|
|
df_dy = 200 * (y - x ^ 2)
|
|
return Float64[df_dx, df_dy]
|
|
end
|
|
|
|
|
|
function optimize(opt::Optimizer, f, grad_f, start::Vector{Float64}; steps::Int=5000)
|
|
params = copy(start)
|
|
history = Vector{Vector{Float64}}()
|
|
push!(history, copy(params))
|
|
for _ in 1:steps
|
|
grads = grad_f(params)
|
|
if any(g -> !isfinite(g) || abs(g) > 1e15, grads)
|
|
break
|
|
end
|
|
params = step!(opt, params, grads)
|
|
if any(p -> !isfinite(p) || abs(p) > 1e15, params)
|
|
break
|
|
end
|
|
push!(history, copy(params))
|
|
end
|
|
return history
|
|
end
|
|
|
|
|
|
function distance_to_minimum(p::Vector{Float64}, target::Tuple{Float64, Float64}=(1.0, 1.0))::Float64
|
|
return sqrt((p[1] - target[1]) ^ 2 + (p[2] - target[2]) ^ 2)
|
|
end
|
|
|
|
|
|
function find_convergence_step(history, f; threshold::Float64=1e-4)::Int
|
|
for (i, params) in enumerate(history)
|
|
if f(params) < threshold
|
|
return i - 1
|
|
end
|
|
end
|
|
return length(history)
|
|
end
|
|
|
|
|
|
function print_trajectory(name::String, history, f; steps_to_show::Int=10)
|
|
total = length(history) - 1
|
|
interval = max(1, total ÷ steps_to_show)
|
|
println("\n" * "=" ^ 60)
|
|
println(" $name")
|
|
println("=" ^ 60)
|
|
@printf(" %6s %10s %10s %14s %8s\n", "Step", "x", "y", "Loss", "Dist")
|
|
println(" " * "-" ^ 52)
|
|
for i in 0:interval:total
|
|
p = history[i + 1]
|
|
loss = f(p)
|
|
dist = distance_to_minimum(p)
|
|
@printf(" %6d %10.6f %10.6f %14.8f %8.4f\n", i, p[1], p[2], loss, dist)
|
|
end
|
|
if total % interval != 0
|
|
p = history[end]
|
|
loss = f(p)
|
|
dist = distance_to_minimum(p)
|
|
@printf(" %6d %10.6f %10.6f %14.8f %8.4f\n", total, p[1], p[2], loss, dist)
|
|
end
|
|
end
|
|
|
|
|
|
function print_ascii_convergence(results, f; steps::Int=5000)
|
|
println("\n" * "=" ^ 60)
|
|
println(" CONVERGENCE COMPARISON (log10 loss over steps)")
|
|
println("=" ^ 60)
|
|
width = 50
|
|
sample_points = 40
|
|
interval = max(1, steps ÷ sample_points)
|
|
for (name, history) in results
|
|
losses = Float64[]
|
|
i = 0
|
|
while i <= min(length(history) - 1, steps)
|
|
push!(losses, f(history[i + 1]))
|
|
i += interval
|
|
end
|
|
isempty(losses) && continue
|
|
max_log = 5.0
|
|
min_log = -8.0
|
|
log_range = max_log - min_log
|
|
bars = Int[]
|
|
for loss in losses
|
|
ll = log10(loss + 1e-15)
|
|
ll = clamp(ll, min_log, max_log)
|
|
normalized = (ll - min_log) / log_range
|
|
push!(bars, Int(round(normalized * (width - 1))))
|
|
end
|
|
println("\n $name:")
|
|
println(" loss 1e-8 " * "."^width * " 1e+5")
|
|
for (idx, pos) in enumerate(bars)
|
|
step_num = (idx - 1) * interval
|
|
line = fill(' ', width)
|
|
line[clamp(pos + 1, 1, width)] = '*'
|
|
println(" " * lpad(string(step_num), 5) * " |" * String(line) * "|")
|
|
end
|
|
final_loss = f(history[end])
|
|
conv_step = find_convergence_step(history, f)
|
|
conv_msg = conv_step < length(history) ? "step $conv_step" : "did not converge"
|
|
@printf(" final loss: %.2e, converged (< 1e-4): %s\n", final_loss, conv_msg)
|
|
end
|
|
end
|
|
|
|
|
|
function demo_comparison()
|
|
println("OPTIMIZATION METHODS COMPARISON")
|
|
println("Minimizing the Rosenbrock function: f(x, y) = (1-x)^2 + 100(y-x^2)^2")
|
|
println("Global minimum at (1, 1) where f = 0")
|
|
@printf("Starting point: (-1.0, 1.0), f = %.1f\n", rosenbrock(Float64[-1.0, 1.0]))
|
|
|
|
start = Float64[-1.0, 1.0]
|
|
steps = 5000
|
|
|
|
configs = [
|
|
("Gradient Descent", GradientDescent(lr=0.0005)),
|
|
("SGD + Momentum", SGDMomentum(lr=0.0001, momentum=0.9)),
|
|
("Adam", Adam(lr=0.01)),
|
|
]
|
|
|
|
results = Tuple{String, Vector{Vector{Float64}}}[]
|
|
for (name, opt) in configs
|
|
history = optimize(opt, rosenbrock, rosenbrock_grad, start; steps=steps)
|
|
push!(results, (name, history))
|
|
print_trajectory(name, history, rosenbrock)
|
|
end
|
|
|
|
print_ascii_convergence(results, rosenbrock; steps=steps)
|
|
|
|
println("\n" * "=" ^ 60)
|
|
println(" FINAL RESULTS")
|
|
println("=" ^ 60)
|
|
@printf(" %-22s %10s %10s %14s\n", "Method", "x", "y", "Loss")
|
|
println(" " * "-" ^ 58)
|
|
for (name, history) in results
|
|
final = history[end]
|
|
loss = rosenbrock(final)
|
|
@printf(" %-22s %10.6f %10.6f %14.8f\n", name, final[1], final[2], loss)
|
|
end
|
|
println("\n Target: x=1.000000, y=1.000000, loss=0.00000000")
|
|
end
|
|
|
|
|
|
function demo_learning_rate_effect()
|
|
println("\n\n" * "=" ^ 60)
|
|
println(" LEARNING RATE EFFECT ON GRADIENT DESCENT")
|
|
println("=" ^ 60)
|
|
start = Float64[-1.0, 1.0]
|
|
rates = [0.0001, 0.0005, 0.001, 0.005]
|
|
@printf("\n %8s %10s %10s %14s %s\n", "LR", "Final x", "Final y", "Loss", "Status")
|
|
println(" " * "-" ^ 60)
|
|
for lr in rates
|
|
gd = GradientDescent(lr=lr)
|
|
history = optimize(gd, rosenbrock, rosenbrock_grad, start; steps=5000)
|
|
final = history[end]
|
|
loss = rosenbrock(final)
|
|
diverged = !isfinite(loss) || loss > 1e10
|
|
status = diverged ? "DIVERGED" : (loss < 0.01 ? "converged" : "slow")
|
|
if diverged
|
|
@printf(" %8.4f %10s %10s %14s %s\n", lr, "nan", "nan", "inf", status)
|
|
else
|
|
@printf(" %8.4f %10.6f %10.6f %14.8f %s\n", lr, final[1], final[2], loss, status)
|
|
end
|
|
end
|
|
end
|
|
|
|
|
|
function demo_momentum_effect()
|
|
println("\n\n" * "=" ^ 60)
|
|
println(" MOMENTUM EFFECT ON SGD")
|
|
println("=" ^ 60)
|
|
start = Float64[-1.0, 1.0]
|
|
betas = [0.0, 0.5, 0.9, 0.99]
|
|
@printf("\n %6s %10s %10s %14s\n", "Beta", "Final x", "Final y", "Loss")
|
|
println(" " * "-" ^ 46)
|
|
for beta in betas
|
|
sgd = SGDMomentum(lr=0.0001, momentum=beta)
|
|
history = optimize(sgd, rosenbrock, rosenbrock_grad, start; steps=5000)
|
|
final = history[end]
|
|
loss = rosenbrock(final)
|
|
if !isfinite(loss)
|
|
@printf(" %6.2f %10s %10s %14s\n", beta, "nan", "nan", "inf")
|
|
else
|
|
@printf(" %6.2f %10.6f %10.6f %14.8f\n", beta, final[1], final[2], loss)
|
|
end
|
|
end
|
|
end
|
|
|
|
|
|
function demo_saddle_point()
|
|
println("\n\n" * "=" ^ 60)
|
|
println(" SADDLE POINT ESCAPE: f(x, y) = x^2 - y^2")
|
|
println("=" ^ 60)
|
|
|
|
saddle(p::Vector{Float64}) = p[1] ^ 2 - p[2] ^ 2
|
|
saddle_grad(p::Vector{Float64}) = Float64[2 * p[1], -2 * p[2]]
|
|
|
|
start = Float64[0.01, 0.01]
|
|
steps = 200
|
|
|
|
configs = [
|
|
("Gradient Descent", GradientDescent(lr=0.01)),
|
|
("SGD + Momentum", SGDMomentum(lr=0.01, momentum=0.9)),
|
|
("Adam", Adam(lr=0.01)),
|
|
]
|
|
|
|
println("\n Start: x=0.01, y=0.01 (near saddle at origin)")
|
|
@printf("\n %-22s %10s %10s %12s %s\n", "Method", "x", "y", "f(x, y)", "Escaped?")
|
|
println(" " * "-" ^ 62)
|
|
for (name, opt) in configs
|
|
history = optimize(opt, saddle, saddle_grad, start; steps=steps)
|
|
final = history[end]
|
|
val = saddle(final)
|
|
escaped = abs(final[2]) > 1.0 ? "yes" : "no"
|
|
@printf(" %-22s %10.6f %10.6f %12.6f %s\n", name, final[1], final[2], val, escaped)
|
|
end
|
|
end
|
|
|
|
|
|
function main()
|
|
demo_comparison()
|
|
demo_learning_rate_effect()
|
|
demo_momentum_effect()
|
|
demo_saddle_point()
|
|
end
|
|
|
|
|
|
if abspath(PROGRAM_FILE) == @__FILE__
|
|
main()
|
|
end
|