Skip to content

Instantly share code, notes, and snippets.

@torfjelde
Last active February 27, 2024 21:19
Show Gist options
  • Save torfjelde/1fe3a149ead219103d6b535217c7d2d9 to your computer and use it in GitHub Desktop.
Save torfjelde/1fe3a149ead219103d6b535217c7d2d9 to your computer and use it in GitHub Desktop.
Determining chunksize to use in ForwardDiff.jl for a Turing.jl model
julia> using Turing, TuringBenchmarking, ADTypes
julia> @model function f_model(n)
μ ~ truncated(Normal(), lower=0)
σ² ~ InverseGamma(2, 3)
x ~ filldist(truncated(Normal(μ, sqrt(σ²)); lower=0), n)
end
f_model (generic function with 2 methods)
julia> max_chunksizes = [
0, # corresponds to the default heuristic
1,
4,
8,
16,
32
];
julia> # Turing.jl now uses ADTypes.jl for the AD backend configuration.
adbackends=map(c -> ADTypes.AutoForwardDiff(chunksize=c), chunksizes)
6-element Vector{AutoForwardDiff{_A, Nothing} where _A}:
AutoForwardDiff{0, Nothing}(nothing)
AutoForwardDiff{1, Nothing}(nothing)
AutoForwardDiff{4, Nothing}(nothing)
AutoForwardDiff{8, Nothing}(nothing)
AutoForwardDiff{16, Nothing}(nothing)
AutoForwardDiff{32, Nothing}(nothing)
julia> # Benchmark the model with different backends.
results = TuringBenchmarking.benchmark_model(f_model(100); adbackends)
2-element BenchmarkTools.BenchmarkGroup:
tags: []
"evaluation" => 2-element BenchmarkTools.BenchmarkGroup:
tags: []
"linked" => Trial(4.727 μs)
"standard" => Trial(1.906 μs)
"gradient" => 6-element BenchmarkTools.BenchmarkGroup:
tags: []
"AutoForwardDiff{4, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(251.729 μs)
"standard" => Trial(123.806 μs)
"AutoForwardDiff{32, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(89.686 μs)
"standard" => Trial(46.526 μs)
"AutoForwardDiff{8, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(141.020 μs)
"standard" => Trial(74.930 μs)
"AutoForwardDiff{16, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(93.164 μs)
"standard" => Trial(46.667 μs)
"AutoForwardDiff{1, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(635.984 μs)
"standard" => Trial(241.602 μs)
"AutoForwardDiff{0, Nothing}(nothing)" => 2-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(110.319 μs)
"standard" => Trial(56.452 μs)
julia> # Extract the results for the gradient computation with linking.
results_gradient = results[@tagged ("gradient" && "linked")]
1-element BenchmarkTools.BenchmarkGroup:
tags: []
"gradient" => 6-element BenchmarkTools.BenchmarkGroup:
tags: []
"AutoForwardDiff{4, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(251.729 μs)
"AutoForwardDiff{32, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(89.686 μs)
"AutoForwardDiff{8, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(141.020 μs)
"AutoForwardDiff{16, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(93.164 μs)
"AutoForwardDiff{1, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(635.984 μs)
"AutoForwardDiff{0, Nothing}(nothing)" => 1-element BenchmarkTools.BenchmarkGroup:
tags: ["ForwardDiff"]
"linked" => Trial(110.319 μs)
julia> # Sort according to the minimum time for each backend.
sort(
BenchmarkTools.leaves(results_gradient),
by=minimum ∘ last
)
6-element Vector{Any}:
(Any["gradient", "AutoForwardDiff{32, Nothing}(nothing)", "linked"], Trial(89.686 μs))
(Any["gradient", "AutoForwardDiff{16, Nothing}(nothing)", "linked"], Trial(93.164 μs))
(Any["gradient", "AutoForwardDiff{0, Nothing}(nothing)", "linked"], Trial(110.319 μs))
(Any["gradient", "AutoForwardDiff{8, Nothing}(nothing)", "linked"], Trial(141.020 μs))
(Any["gradient", "AutoForwardDiff{4, Nothing}(nothing)", "linked"], Trial(251.729 μs))
(Any["gradient", "AutoForwardDiff{1, Nothing}(nothing)", "linked"], Trial(635.984 μs))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment