Update models to benchmark plus small style changes

Author: mhauru

benchmarks/benchmarks.jl

@@ -1,39 +1,46 @@
-using DynamicPPL: @model
-using DynamicPPLBenchmarks: make_suite
+using DynamicPPLBenchmarks: Models, make_suite
 using BenchmarkTools: median, run
-using Distributions: Normal, Beta, Bernoulli
-using PrettyTables: pretty_table, PrettyTables
+using PrettyTables: PrettyTables
+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)
+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)
 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),
+    ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff),
+    ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff),
+    ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff),
+    ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff),
+    ("Smorgasbord", smorgasbord_instance, :typed, :reversediff),
+    #("Smorgasbord", smorgasbord_instance, :typed, :mooncake),
+    ("Loop univariate 1k", loop_univariate1k, :typed, :reversediff),
+    ("Multivariate 1k", multivariate1k, :typed, :reversediff),
+    ("Loop univariate 10k", loop_univariate10k, :typed, :reversediff),
+    ("Multivariate 10k", multivariate10k, :typed, :reversediff),
+    ("Dynamic", Models.dynamic(), :typed, :reversediff),
+    ("Submodel", Models.parent(randn()), :typed, :reversediff),
+    ("LDA", lda_instance, :typed, :reversediff),
 ]
 
 results_table = Tuple{String,String,String,Float64,Float64}[]
@@ -42,9 +49,9 @@ for (model_name, model, varinfo_choice, adbackend) in chosen_combinations
     suite = make_suite(model, varinfo_choice, adbackend)
     results = run(suite)
 
-    eval_time = median(results["AD_Benchmarking"]["evaluation"]["standard"]).time
+    eval_time = median(results["evaluation"]["standard"]).time
 
-    grad_group = results["AD_Benchmarking"]["gradient"]
+    grad_group = results["gradient"]
     if isempty(grad_group)
         ad_eval_time = NaN
     else
@@ -62,4 +69,4 @@ table_matrix = hcat(Iterators.map(collect, zip(results_table...))...)
 header = [
     "Model", "AD Backend", "VarInfo Type", "Evaluation Time (ns)", "AD Eval Time (ns)"
 ]
-pretty_table(table_matrix; header=header, tf=PrettyTables.tf_markdown)
+PrettyTables.pretty_table(table_matrix; header=header, tf=PrettyTables.tf_markdown)

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,121 @@
+"""
+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 DynamicPPL: @model, to_submodel
+using Distributions:
+    Categorical,
+    Dirichlet,
+    Exponential,
+    Gamma,
+    LKJCholesky,
+    MatrixBeta,
+    Normal,
+    logpdf,
+    product_distribution,
+    truncated
+
+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(x, obs)
+    logp = logdf(x, Normal())
+    logp += logpdf(obs, Normal(x, 1))
+    return logp
+end
+
+@model function simple_assume_observe(obs)
+    x ~ Normal()
+    obs ~ Normal(x, 1)
+    return (; x=x)
+end
+
+@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
+
+@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
+
+@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
+
+@model function sub()
+    x ~ Normal()
+    return x
+end
+
+@model function parent(y)
+    x ~ to_submodel(sub())
+    y ~ Normal(x, 1)
+    return (; x=x)
+end
+
+@model function dynamic(::Type{T}=Vector{Float64}) where {T}
+    eta ~ truncated(Normal(); lower=0.0)
+    mat1 ~ LKJCholesky(4, eta)
+    mat2 ~ MatrixBeta(5, 6.0, 8.0)
+    dim = eta > 0.2 ? 2 : 3
+    vec = T(undef, dim)
+    vec .~ truncated(Exponential(0.5); lower=0.0, upper=1.0)
+    return (; eta=eta, mat1=mat1, mat2=mat2, vec=vec)
+end
+
+@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