diff --git a/HISTORY.md b/HISTORY.md index b59d8dd7f..7a581e0cf 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -24,9 +24,12 @@ Please see the API documentation for more details. There is now also an `rng` keyword argument to help seed parameter generation. -Finally, instead of specifying `value_atol` and `grad_atol`, you can now specify `atol` and `rtol` which are used for both value and gradient. +Instead of specifying `value_atol` and `grad_atol`, you can now specify `atol` and `rtol` which are used for both value and gradient. Their semantics are the same as in Julia's `isapprox`; two values are equal if they satisfy either `atol` or `rtol`. +Finally, the `ADResult` object returned by `run_ad` now has both `grad_time` and `primal_time` fields, which contain the time it took to calculate the gradient of logp and logp itself. +Previously there was only a single `time_vs_primal` field which represented the ratio of these two. + ### `DynamicPPL.TestUtils.check_model` You now need to explicitly pass a `VarInfo` argument to `check_model` and `check_model_and_trace`. diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 3d14d03ff..137544e56 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -4,28 +4,28 @@ version = "0.1.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [sources] DynamicPPL = {path = "../"} [compat] -ADTypes = "1.14.0" -BenchmarkTools = "1.6.0" -Distributions = "0.25.117" +ADTypes = "1" +Chairmarks = "1" +Distributions = "0.25" DynamicPPL = "0.37" ForwardDiff = "0.10.38, 1" -LogDensityProblems = "2.1.2" Mooncake = "0.4" -PrettyTables = "2.4.0" -ReverseDiff = "1.15.3" +PrettyTables = "2" +ReverseDiff = "1" StableRNGs = "1" +Statistics = "1" diff --git a/benchmarks/benchmarks.jl b/benchmarks/benchmarks.jl index b733d810c..47a8a2912 100644 --- a/benchmarks/benchmarks.jl +++ b/benchmarks/benchmarks.jl @@ -1,9 +1,9 @@ -using Pkg - -using DynamicPPLBenchmarks: Models, make_suite, model_dimension -using BenchmarkTools: @benchmark, median, run +using DynamicPPLBenchmarks: Models, to_backend, make_varinfo +using DynamicPPL.TestUtils.AD: run_ad, NoTest +using Chairmarks: @be using PrettyTables: PrettyTables, ft_printf using StableRNGs: StableRNG +using Statistics: median rng = StableRNG(23) @@ -35,48 +35,45 @@ chosen_combinations = [ Models.simple_assume_observe(randn(rng)), :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), - ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true), - ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true), - ("Multivariate 1k", multivariate1k, :typed, :mooncake, true), - ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true), - ("Multivariate 10k", multivariate10k, :typed, :mooncake, true), - ("Dynamic", Models.dynamic(), :typed, :mooncake, true), - ("Submodel", Models.parent(randn(rng)), :typed, :mooncake, true), - ("LDA", lda_instance, :typed, :reversediff, true), + ("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, :mooncake), + ("Multivariate 1k", multivariate1k, :typed, :mooncake), + ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake), + ("Multivariate 10k", multivariate10k, :typed, :mooncake), + ("Dynamic", Models.dynamic(), :typed, :mooncake), + ("Submodel", Models.parent(randn(rng)), :typed, :mooncake), + ("LDA", lda_instance, :typed, :reversediff), ] # 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(rng) - median(@benchmark Models.simple_assume_observe_non_model(obs)).time + median(@be Models.simple_assume_observe_non_model(obs)).time end -results_table = Tuple{String,Int,String,String,Bool,Float64,Float64}[] +results_table = Tuple{String,Int,String,String,Float64,Float64}[] -for (model_name, model, varinfo_choice, adbackend, islinked) in chosen_combinations +for (model_name, model, varinfo_choice, adbackend) in chosen_combinations @info "Running benchmark for $model_name" - suite = make_suite(model, varinfo_choice, adbackend, islinked) - results = run(suite) - eval_time = median(results["evaluation"]).time - relative_eval_time = eval_time / reference_time - ad_eval_time = median(results["gradient"]).time - relative_ad_eval_time = ad_eval_time / eval_time + adtype = to_backend(adbackend) + varinfo = make_varinfo(model, varinfo_choice) + ad_result = run_ad(model, adtype; test=NoTest(), benchmark=true, varinfo=varinfo) + relative_eval_time = ad_result.primal_time / reference_time + relative_ad_eval_time = ad_result.grad_time / ad_result.primal_time push!( results_table, ( model_name, - model_dimension(model, islinked), + length(varinfo[:]), string(adbackend), string(varinfo_choice), - islinked, relative_eval_time, relative_ad_eval_time, ), @@ -89,7 +86,6 @@ header = [ "Dimension", "AD Backend", "VarInfo Type", - "Linked", "Eval Time / Ref Time", "AD Time / Eval Time", ] @@ -97,6 +93,6 @@ PrettyTables.pretty_table( table_matrix; header=header, tf=PrettyTables.tf_markdown, - formatters=ft_printf("%.1f", [6, 7]), + formatters=ft_printf("%.1f", [5, 6]), crop=:none, # Always print the whole table, even if it doesn't fit in the terminal. ) diff --git a/benchmarks/src/DynamicPPLBenchmarks.jl b/benchmarks/src/DynamicPPLBenchmarks.jl index 8c5032ace..d204c1c0d 100644 --- a/benchmarks/src/DynamicPPLBenchmarks.jl +++ b/benchmarks/src/DynamicPPLBenchmarks.jl @@ -1,10 +1,8 @@ module DynamicPPLBenchmarks -using DynamicPPL: VarInfo, SimpleVarInfo, VarName -using BenchmarkTools: BenchmarkGroup, @benchmarkable +using DynamicPPL: Model, VarInfo, SimpleVarInfo using DynamicPPL: DynamicPPL using ADTypes: ADTypes -using LogDensityProblems: LogDensityProblems using ForwardDiff: ForwardDiff using Mooncake: Mooncake @@ -14,29 +12,14 @@ using StableRNGs: StableRNG include("./Models.jl") using .Models: Models -export Models, make_suite, model_dimension - -""" - model_dimension(model, islinked) - -Return the dimension of `model`, accounting for linking, if any. -""" -function model_dimension(model, islinked) - vi = VarInfo() - model(StableRNG(23), vi) - if islinked - vi = DynamicPPL.link(vi, model) - end - return length(vi[:]) -end +export Models, to_backend, make_varinfo # Utility functions for representing AD backends using symbols. # Copied from TuringBenchmarking.jl. const SYMBOL_TO_BACKEND = Dict( :forwarddiff => ADTypes.AutoForwardDiff(), :reversediff => ADTypes.AutoReverseDiff(; compile=false), - :reversediff_compiled => ADTypes.AutoReverseDiff(; compile=true), - :mooncake => ADTypes.AutoMooncake(; config=nothing), + :mooncake => ADTypes.AutoMooncake(), ) to_backend(x) = error("Unknown backend: $x") @@ -48,58 +31,37 @@ function to_backend(x::Union{AbstractString,Symbol}) end """ - make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool) + make_varinfo(model, varinfo_choice::Symbol) -Create a benchmark suite for `model` using the selected varinfo type and AD backend. +Create a VarInfo for the given `model` using the selected varinfo type. Available varinfo choices: • `:untyped` → uses `DynamicPPL.untyped_varinfo(model)` • `:typed` → uses `DynamicPPL.typed_varinfo(model)` - • `: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`). + • `:simple_namedtuple` → builds a `SimpleVarInfo{Float64}(::NamedTuple)` + • `:simple_dict` → builds a `SimpleVarInfo{Float64}(::Dict)` -`islinked` determines whether to link the VarInfo for evaluation. +The VarInfo is always linked. """ -function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::Bool) +function make_varinfo(model::Model, varinfo_choice::Symbol) rng = StableRNG(23) - suite = BenchmarkGroup() - vi = if varinfo_choice == :untyped DynamicPPL.untyped_varinfo(rng, model) elseif varinfo_choice == :typed DynamicPPL.typed_varinfo(rng, model) elseif varinfo_choice == :simple_namedtuple - SimpleVarInfo{Float64}(model(rng)) + vi = DynamicPPL.typed_varinfo(rng, model) + vals = DynamicPPL.values_as(vi, NamedTuple) + SimpleVarInfo{Float64}(vals) elseif varinfo_choice == :simple_dict - retvals = model(rng) - vns = [VarName{k}() for k in keys(retvals)] - SimpleVarInfo{Float64}(Dict(zip(vns, values(retvals)))) + vi = DynamicPPL.typed_varinfo(rng, model) + vals = DynamicPPL.values_as(vi, Dict) + SimpleVarInfo{Float64}(vals) else error("Unknown varinfo choice: $varinfo_choice") end - adbackend = to_backend(adbackend) - - if islinked - vi = DynamicPPL.link(vi, model) - end - - f = DynamicPPL.LogDensityFunction( - model, DynamicPPL.getlogjoint_internal, vi; adtype=adbackend - ) - # The parameters at which we evaluate f. - θ = vi[:] - - # Run once to trigger compilation. - LogDensityProblems.logdensity_and_gradient(f, θ) - suite["gradient"] = @benchmarkable $(LogDensityProblems.logdensity_and_gradient)($f, $θ) - - # Also benchmark just standard model evaluation because why not. - suite["evaluation"] = @benchmarkable $(LogDensityProblems.logdensity)($f, $θ) - - return suite + return DynamicPPL.link!!(vi, model) end end # module diff --git a/src/test_utils/ad.jl b/src/test_utils/ad.jl index 1ac33a481..068e1f346 100644 --- a/src/test_utils/ad.jl +++ b/src/test_utils/ad.jl @@ -109,8 +109,11 @@ struct ADResult{Tparams<:AbstractFloat,Tresult<:AbstractFloat,Ttol<:AbstractFloa value_actual::Tresult "The gradient of logp (calculated using `adtype`)" grad_actual::Vector{Tresult} - "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself" - time_vs_primal::Union{Nothing,Tresult} + "If benchmarking was requested, the time taken by the AD backend to evaluate the gradient + of logp" + grad_time::Union{Nothing,Tresult} + "If benchmarking was requested, the time taken by the AD backend to evaluate logp" + primal_time::Union{Nothing,Tresult} end """ @@ -121,6 +124,8 @@ end benchmark=false, atol::AbstractFloat=1e-8, rtol::AbstractFloat=sqrt(eps()), + getlogdensity::Function=getlogjoint_internal, + rng::AbstractRNG=default_rng(), varinfo::AbstractVarInfo=link(VarInfo(model), model), params::Union{Nothing,Vector{<:AbstractFloat}}=nothing, verbose=true, @@ -174,6 +179,21 @@ Everything else is optional, and can be categorised into several groups: prep_params)`. You could then evaluate the gradient at a different set of parameters using the `params` keyword argument. +3. _Which type of logp is being calculated._ + + By default, `run_ad` evaluates the 'internal log joint density' of the model, + i.e., the log joint density in the unconstrained space. Thus, for example, in + + @model f() = x ~ LogNormal() + + the internal log joint density is `logpdf(Normal(), log(x))`. This is the + relevant log density for e.g. Hamiltonian Monte Carlo samplers and is therefore + the most useful to test. + + If you want the log joint density in the original model parameterisation, you + can use `getlogjoint`. Likewise, if you want only the prior or likelihood, + you can use `getlogprior` or `getloglikelihood`, respectively. + 3. _How to specify the results to compare against._ Once logp and its gradient has been calculated with the specified `adtype`, @@ -277,14 +297,18 @@ function run_ad( end # Benchmark - time_vs_primal = if benchmark + grad_time, primal_time = if benchmark primal_benchmark = @be (ldf, params) logdensity(_[1], _[2]) grad_benchmark = @be (ldf, params) logdensity_and_gradient(_[1], _[2]) - t = median(grad_benchmark).time / median(primal_benchmark).time - verbose && println("grad / primal : $(t)") - t + median_primal = median(primal_benchmark).time + median_grad = median(grad_benchmark).time + r(f) = round(f; sigdigits=4) + verbose && println( + "grad / primal : $(r(median_grad))/$(r(median_primal)) = $(r(median_grad / median_primal))", + ) + (median_grad, median_primal) else - nothing + nothing, nothing end return ADResult( @@ -299,7 +323,8 @@ function run_ad( grad_true, value, grad, - time_vs_primal, + grad_time, + primal_time, ) end