|
import Pkg; Pkg.status() |
|
|
|
@info "Loading packages..." |
|
import Random, DelimitedFiles |
|
import ForwardDiff, ReverseDiff, Zygote, Symbolics |
|
using BenchmarkTools |
|
|
|
const AV = AbstractVector{T} where T |
|
|
|
# ========== Objective function ========== |
|
|
|
normal_pdf(x::Real, mean::Real, var::Real) = |
|
exp(-(x - mean)^2 / (2var)) / sqrt(2π * var) |
|
|
|
mixture_pdf(x::Real, weights::AV{<:Real}, means::AV{<:Real}, vars::AV{<:Real}) = |
|
sum( |
|
w * normal_pdf(x, mean, var) |
|
for (w, mean, var) in zip(weights, means, vars) |
|
) |
|
|
|
normal_pdf(x, mean, var) = |
|
exp(-(x - mean)^2 / (2var)) / sqrt(2π * var) |
|
|
|
|
|
function mixture_loglikelihood(params::AV{<:Real}, data::AV{<:Real})::Real |
|
K = length(params) ÷ 3 |
|
weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end] |
|
|
|
mat = normal_pdf.(data, means', stds' .^2) # (N, K) |
|
#@show size(mat) |
|
sum(mat .* weights', dims=2) .|> log |> sum |
|
# sum( |
|
# sum( |
|
# weight * normal_pdf(x, mean, std^2) |
|
# for (weight, mean, std) in zip(weights, means, stds) |
|
# ) |> log |
|
# for x in data |
|
# ) |
|
end |
|
|
|
function generate_gradient(out_fname::AbstractString, K::Integer) |
|
@assert K > 0 |
|
Symbolics.@variables x ws[1:K] mus[1:K] stds[1:K] |
|
|
|
args=[x, ws, mus, stds] |
|
expr = Symbolics.gradient( |
|
mixture_pdf(x, ws, mus, collect(stds) .^2) |> log, |
|
[ws; mus; stds] |
|
) |
|
|
|
fn, fn_mut = Symbolics.build_function(expr, args...) |
|
|
|
write(out_fname, string(fn_mut)) |
|
end |
|
|
|
# ========== Gradient with Symbolics.jl ========== |
|
@info "Generating gradient functions..." |
|
GRAD_FNS = Union{Nothing, Function}[nothing] |
|
for K in 2:5 |
|
fname = "grad_$K.jl" |
|
@show generate_gradient(fname, K) |
|
push!(GRAD_FNS, include(fname)) |
|
end |
|
|
|
function my_gradient!(out::AV{<:Real}, tmp::AV{<:Real}, xs::AV{<:Real}, params::AV{<:Real}) |
|
K = length(params) ÷ 3 |
|
grad! = GRAD_FNS[K] |
|
weights, means, stds = @views params[1:K], params[K+1:2K], params[2K+1:end] |
|
|
|
out .= 0 |
|
for x in xs |
|
grad!(tmp, x, weights, means, stds) |
|
out .+= tmp |
|
end |
|
end |
|
|
|
|
|
# ========== Benchmark setup ========== |
|
SEED = 42 |
|
N_SAMPLES = 500 |
|
N_COMPONENTS = 4 |
|
|
|
rnd = Random.MersenneTwister(SEED) |
|
data = randn(rnd, N_SAMPLES) |
|
params0 = [rand(rnd, N_COMPONENTS); randn(rnd, N_COMPONENTS); 2rand(rnd, N_COMPONENTS)] |
|
DelimitedFiles.writedlm("gen_data.csv", data, ',') |
|
DelimitedFiles.writedlm("gen_params0.csv", params0, ',') |
|
|
|
objective = params -> mixture_loglikelihood(params, data) |
|
|
|
@show params0 |
|
@show objective(params0) |
|
|
|
@info "Settings" SEED N_SAMPLES N_COMPONENTS length(params0) |
|
|
|
|
|
# ========== Actual benchmarks ========== |
|
@info "Computing gradient w/ Symbolics" |
|
let |
|
grad_storage = similar(params0) |
|
tmp = similar(params0) |
|
|
|
# 1. Compile |
|
my_gradient!(grad_storage, tmp, data, params0) |
|
# 2. Benchmark |
|
trial = run(@benchmarkable $my_gradient!($grad_storage, $tmp, $data, $params0) samples=10_000 evals=1 seconds=60) |
|
show(stdout, MIME("text/plain"), trial) |
|
println() |
|
@show grad_storage |
|
end |
|
|
|
@info "Computing gradient w/ ForwardDiff" |
|
let |
|
grad_storage = similar(params0) |
|
cfg_grad = ForwardDiff.GradientConfig(objective, params0, ForwardDiff.Chunk{length(params0)}()) |
|
|
|
# 1. Compile |
|
ForwardDiff.gradient!(grad_storage, objective, params0, cfg_grad) |
|
# 2. Benchmark |
|
trial = run(@benchmarkable ForwardDiff.gradient!($grad_storage, $objective, $params0, $cfg_grad) samples=10_000 evals=1 seconds=60) |
|
show(stdout, MIME("text/plain"), trial) |
|
println() |
|
@show grad_storage |
|
end |
|
|
|
@info "Computing gradient w/ ReverseDiff" |
|
let |
|
grad_storage = similar(params0) |
|
objective_tape = ReverseDiff.GradientTape(objective, params0) |> ReverseDiff.compile |
|
|
|
# 1. Compile |
|
ReverseDiff.gradient!(grad_storage, objective_tape, params0) |
|
# 2. Benchmark |
|
trial = run(@benchmarkable ReverseDiff.gradient!($grad_storage, $objective_tape, $params0) samples=10_000 evals=1 seconds=60) |
|
show(stdout, MIME("text/plain"), trial) |
|
println() |
|
@show grad_storage |
|
end |
|
|
|
@info "Computing gradient w/ Zygote reverse" |
|
let |
|
# 1. Compile |
|
grad_storage = Zygote.gradient(objective, params0) |
|
# 2. Benchmark |
|
trial = run(@benchmarkable Zygote.gradient($objective, $params0) samples=10_000 evals=1 seconds=60) |
|
show(stdout, MIME("text/plain"), trial) |
|
println() |
|
@show grad_storage |
|
end |
