Skip to content

Commit

Permalink
Insane memory-usage when using Emcee with large number of iterations (
Browse files Browse the repository at this point in the history
#1976)

* use `reduce` instead of splatting to avoid stack overflow

* added test for mem usage of Emcee

* version bump

* for emcee, walkers are now indexed along the chain index rather than
iteration index

* reduced number of iterations for emcee to see if that fixes tests

* attempt at reducing the memory usage of Emcee tests

* attempt at further reducing the memory footprint of the test

* fixed size issues

* Update emcee.jl

---------

Co-authored-by: Hong Ge <[email protected]>
  • Loading branch information
torfjelde and yebai authored Jun 14, 2023
1 parent f8450d6 commit a38c709
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 6 deletions.
9 changes: 7 additions & 2 deletions src/inference/emcee.jl
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,13 @@ function AbstractMCMC.bundle_samples(

# Extract names & construct param array.
nms = [nms; extra_params]
parray = map(x -> hcat(x[1], x[2]), zip(vals_vec, extra_values_vec))
parray = cat(parray..., dims=3)
# `hcat` first to ensure we get the right `eltype`.
x = hcat(first(vals_vec), first(extra_values_vec))
# Pre-allocate to minimize memory usage.
parray = Array{eltype(x),3}(undef, length(vals_vec), size(x, 2), size(x, 1))
for (i, (vals, extras)) in enumerate(zip(vals_vec, extra_values_vec))
parray[i, :, :] = transpose(hcat(vals, extras))
end

# Get the average or final log evidence, if it exists.
le = getlogevidence(samples, state, spl)
Expand Down
15 changes: 11 additions & 4 deletions test/inference/emcee.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@


@testset "emcee.jl" begin
@testset "gdemo" begin
Random.seed!(9876)
Expand All @@ -14,6 +12,15 @@
check_gdemo(chain)
end

@testset "memory usage with large number of iterations" begin
# https://github.com/TuringLang/Turing.jl/pull/1976
@info "Testing emcee with large number of iterations"
spl = Emcee(10, 2.0)
n_samples = 10_000
chain = sample(gdemo_default, spl, n_samples)
check_gdemo(chain)
end

@testset "initial parameters" begin
nwalkers = 250
spl = Emcee(nwalkers, 2.0)
Expand All @@ -30,7 +37,7 @@

# Initial parameters
chain = sample(gdemo_default, spl, 1; init_params=fill([2.0, 1.0], nwalkers))
@test chain[:s] == fill(2.0, nwalkers, 1)
@test chain[:m] == fill(1.0, nwalkers, 1)
@test chain[:s] == fill(2.0, 1, nwalkers)
@test chain[:m] == fill(1.0, 1, nwalkers)
end
end

0 comments on commit a38c709

Please sign in to comment.