diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 28693ad98..fab8aee23 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -6,5 +6,6 @@ version = "0.1.0" BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f" diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index a095a2d72..359026db6 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -1,65 +1,108 @@ -using DynamicPPL: @model -using DynamicPPLBenchmarks: make_suite -using BenchmarkTools: median, run -using Distributions: Normal, Beta, Bernoulli -using PrettyTables: pretty_table, PrettyTables +using DynamicPPLBenchmarks: Models, make_suite +using BenchmarkTools: @benchmark, median, run +using PrettyTables: PrettyTables, ft_printf +using Random: seed! -# Define models -@model function demo1(x) - m ~ Normal() - x ~ Normal(m, 1) - return (m=m, x=x) -end +seed!(23) -@model function demo2(y) - p ~ Beta(1, 1) - N = length(y) - for n in 1:N - y[n] ~ Bernoulli(p) - end - return (; p) +# Create DynamicPPL.Model instances to run benchmarks on. +smorgasbord_instance = Models.smorgasbord(randn(100), randn(100)) +loop_univariate1k, multivariate1k = begin + data_1k = randn(1_000) + loop = Models.loop_univariate(length(data_1k)) | (; o=data_1k) + multi = Models.multivariate(length(data_1k)) | (; o=data_1k) + loop, multi +end +loop_univariate10k, multivariate10k = begin + data_10k = randn(10_000) + loop = Models.loop_univariate(length(data_10k)) | (; o=data_10k) + multi = Models.multivariate(length(data_10k)) | (; o=data_10k) + loop, multi +end +lda_instance = begin + w = [1, 2, 3, 2, 1, 1] + d = [1, 1, 1, 2, 2, 2] + Models.lda(2, d, w) end - -demo1_data = randn() -demo2_data = rand(Bool, 10) - -# Create model instances with the data -demo1_instance = demo1(demo1_data) -demo2_instance = demo2(demo2_data) # Specify the combinations to test: -# (Model Name, model instance, VarInfo choice, AD backend) +# (Model Name, model instance, VarInfo choice, AD backend, linked) chosen_combinations = [ - ("Demo1", demo1_instance, :typed, :forwarddiff), - ("Demo1", demo1_instance, :simple_namedtuple, :zygote), - ("Demo2", demo2_instance, :untyped, :reversediff), - ("Demo2", demo2_instance, :simple_dict, :forwarddiff), + ( + "Simple assume observe", + Models.simple_assume_observe(randn()), + :typed, + :forwarddiff, + false, + ), + ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff, false), + ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true), + ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true), + ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true), + ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true), + # TODO(mhauru) Add Mooncake once TuringBenchmarking.jl supports it. Consider changing + # all the below :reversediffs to :mooncakes too. + #("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true), + ("Loop univariate 1k", loop_univariate1k, :typed, :reversediff, true), + ("Multivariate 1k", multivariate1k, :typed, :reversediff, true), + ("Loop univariate 10k", loop_univariate10k, :typed, :reversediff, true), + ("Multivariate 10k", multivariate10k, :typed, :reversediff, true), + ("Dynamic", Models.dynamic(), :typed, :reversediff, true), + ("Submodel", Models.parent(randn()), :typed, :reversediff, true), + ("LDA", lda_instance, :typed, :reversediff, true), ] -results_table = Tuple{String,String,String,Float64,Float64}[] +# Time running a model-like function that does not use DynamicPPL, as a reference point. +# Eval timings will be relative to this. +reference_time = begin + obs = randn() + median(@benchmark Models.simple_assume_observe_non_model(obs)).time +end + +results_table = Tuple{String,String,String,Bool,Float64,Float64}[] -for (model_name, model, varinfo_choice, adbackend) in chosen_combinations +for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinations suite = make_suite(model, varinfo_choice, adbackend) results = run(suite) + result_key = islinked ? "linked" : "standard" - eval_time = median(results["AD_Benchmarking"]["evaluation"]["standard"]).time + eval_time = median(results["evaluation"][result_key]).time + relative_eval_time = eval_time / reference_time - grad_group = results["AD_Benchmarking"]["gradient"] + grad_group = results["gradient"] if isempty(grad_group) - ad_eval_time = NaN + relative_ad_eval_time = NaN else grad_backend_key = first(keys(grad_group)) - ad_eval_time = median(grad_group[grad_backend_key]["standard"]).time + ad_eval_time = median(grad_group[grad_backend_key][result_key]).time + relative_ad_eval_time = ad_eval_time / eval_time end push!( results_table, - (model_name, string(adbackend), string(varinfo_choice), eval_time, ad_eval_time), + ( + model_name, + string(adbackend), + string(varinfo_choice), + islinked, + relative_eval_time, + relative_ad_eval_time, + ), ) end table_matrix = hcat(Iterators.map(collect, zip(results_table...))...) header = [ - "Model", "AD Backend", "VarInfo Type", "Evaluation Time (ns)", "AD Eval Time (ns)" + "Model", + "AD Backend", + "VarInfo Type", + "Linked", + "Eval Time / Ref Time", + "AD Time / Eval Time", ] -pretty_table(table_matrix; header=header, tf=PrettyTables.tf_markdown) +PrettyTables.pretty_table( + table_matrix; + header=header, + tf=PrettyTables.tf_markdown, + formatters=ft_printf("%.1f", [5, 6]), +) diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 0b74846b3..f98aede2b 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/benchmarks/src/DynamicPPLBenchmarks.jl @@ -4,7 +4,10 @@ using DynamicPPL: VarInfo, SimpleVarInfo, VarName using BenchmarkTools: BenchmarkGroup using TuringBenchmarking: make_turing_suite -export make_suite +include("./Models.jl") +using .Models: Models + +export Models, make_suite """ make_suite(model, varinfo_choice::Symbol, adbackend::Symbol) @@ -13,7 +16,7 @@ Create a benchmark suite for `model` using the selected varinfo type and AD back Available varinfo choices: • `:untyped` → uses `VarInfo()` • `:typed` → uses `VarInfo(model)` - • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(free_nt)` + • `:simple_namedtuple` → uses `SimpleVarInfo{Float64}(model())` • `:simple_dict` → builds a `SimpleVarInfo{Float64}` from a Dict (pre-populated with the model’s outputs) The AD backend should be specified as a Symbol (e.g. `:forwarddiff`, `:reversediff`, `:zygote`). @@ -22,14 +25,13 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol) suite = BenchmarkGroup() vi = if varinfo_choice == :untyped - v = VarInfo() - model(v) - v + vi = VarInfo() + model(vi) + vi elseif varinfo_choice == :typed VarInfo(model) elseif varinfo_choice == :simple_namedtuple - free_nt = NamedTuple{(:m,)}(model()) # Extract only the free parameter(s) - SimpleVarInfo{Float64}(free_nt) + SimpleVarInfo{Float64}(model()) elseif varinfo_choice == :simple_dict retvals = model() vns = [VarName{k}() for k in keys(retvals)] @@ -39,7 +41,7 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol) end # Add the AD benchmarking suite. - suite["AD_Benchmarking"] = make_turing_suite( + suite = make_turing_suite( model; adbackends=[adbackend], varinfo=vi, diff --git a/benchmarks/src/Models.jl b/benchmarks/src/Models.jl new file mode 100644 index 000000000..25d90f5e1 --- /dev/null +++ b/benchmarks/src/Models.jl @@ -0,0 +1,156 @@ +""" +Models for benchmarking Turing.jl. + +Each model returns a NamedTuple of all the random variables in the model that are not +observed (this is used for constructing SimpleVarInfos). +""" +module Models + +using Distributions: + Categorical, + Dirichlet, + Exponential, + Gamma, + LKJCholesky, + InverseWishart, + Normal, + logpdf, + product_distribution, + truncated +using DynamicPPL: @model, to_submodel +using LinearAlgebra: cholesky + +export simple_assume_observe_non_model, + simple_assume_observe, smorgasbord, loop_univariate, multivariate, parent, dynamic, lda + +# This one is like simple_assume_observe, but explicitly does not use DynamicPPL. +# Other runtimes are normalised by this one's runtime. +function simple_assume_observe_non_model(obs) + x = rand(Normal()) + logp = logpdf(Normal(), x) + logp += logpdf(Normal(x, 1), obs) + return (; logp=logp, x=x) +end + +""" +A simple model that does one scalar assumption and one scalar observation. +""" +@model function simple_assume_observe(obs) + x ~ Normal() + obs ~ Normal(x, 1) + return (; x=x) +end + +""" +A short model that tries to cover many DynamicPPL features. + +Includes scalar, vector univariate, and multivariate variables; ~, .~, and loops; allocating +a variable vector; observations passed as arguments, and as literals. +""" +@model function smorgasbord(x, y, ::Type{TV}=Vector{Float64}) where {TV} + @assert length(x) == length(y) + m ~ truncated(Normal(); lower=0) + means ~ product_distribution(fill(Exponential(m), length(x))) + stds = TV(undef, length(x)) + stds .~ Gamma(1, 1) + for i in 1:length(x) + x[i] ~ Normal(means[i], stds[i]) + end + y ~ product_distribution([Normal(means[i], stds[i]) for i in 1:length(x)]) + 0.0 ~ Normal(sum(y), 1) + return (; m=m, means=means, stds=stds) +end + +""" +A model that loops over two vectors of univariate normals of length `num_dims`. + +The second variable, `o`, is meant to be conditioned on after model instantiation. + +See `multivariate` for a version that uses `product_distribution` rather than loops. +""" +@model function loop_univariate(num_dims, ::Type{TV}=Vector{Float64}) where {TV} + a = TV(undef, num_dims) + o = TV(undef, num_dims) + for i in 1:num_dims + a[i] ~ Normal(0, 1) + end + m = sum(a) + for i in 1:num_dims + o[i] ~ Normal(m, 1) + end + return (; a=a) +end + +""" +A model with two multivariate normal distributed variables of dimension `num_dims`. + +The second variable, `o`, is meant to be conditioned on after model instantiation. + +See `loop_univariate` for a version that uses loops rather than `product_distribution`. +""" +@model function multivariate(num_dims, ::Type{TV}=Vector{Float64}) where {TV} + a = TV(undef, num_dims) + o = TV(undef, num_dims) + a ~ product_distribution(fill(Normal(0, 1), num_dims)) + m = sum(a) + o ~ product_distribution(fill(Normal(m, 1), num_dims)) + return (; a=a) +end + +""" +A submodel for `parent`. Not exported. +""" +@model function sub() + x ~ Normal() + return x +end + +""" +Like simple_assume_observe, but with a submodel for the assumed random variable. +""" +@model function parent(obs) + x ~ to_submodel(sub()) + obs ~ Normal(x, 1) + return (; x=x) +end + +""" +A model with random variables that have changing support under linking, or otherwise +complicated bijectors. +""" +@model function dynamic(::Type{T}=Vector{Float64}) where {T} + eta ~ truncated(Normal(); lower=0.0, upper=0.1) + mat1 ~ LKJCholesky(4, eta) + mat2 ~ InverseWishart(3.2, cholesky([1.0 0.5; 0.5 1.0])) + return (; eta=eta, mat1=mat1, mat2=mat2) +end + +""" +A simple Linear Discriminant Analysis model. +""" +@model function lda(K, d, w) + V = length(unique(w)) + D = length(unique(d)) + N = length(d) + @assert length(w) == N + + ϕ = Vector{Vector{Real}}(undef, K) + for i in 1:K + ϕ[i] ~ Dirichlet(ones(V) / V) + end + + θ = Vector{Vector{Real}}(undef, D) + for i in 1:D + θ[i] ~ Dirichlet(ones(K) / K) + end + + z = zeros(Int, N) + + for i in 1:N + z[i] ~ Categorical(θ[d[i]]) + w[i] ~ Categorical(ϕ[d[i]]) + end + return (; ϕ=ϕ, θ=θ, z=z) +end + +end