Update benchmark models (#826)

* Update models to benchmark plus small style changes

* Make benchmark times relative. Add benchmark documentation.

* Choose whether to show linked or unlinked benchmark times

* Make table header more concise

Author: mhauru

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"

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]),
+)

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,

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