Closed Joshuaalbert closed 8 months ago
For others, they can run the benchmark in benchmarks/run_benchmark.sh
Timings on python 3.11:
2.3.0
Avg. time taken: 7.17404 seconds.
The best 3 of 10 runs took 7.16335 seconds.
2.3.1
Avg. time taken: 7.20859 seconds.
The best 3 of 10 runs took 7.18449 seconds.
2.3.2
Avg. time taken: 7.38937 seconds.
The best 3 of 10 runs took 7.15356 seconds.
2.3.4
Avg. time taken: 7.20582 seconds.
The best 3 of 10 runs took 7.16257 seconds.
2.4.0
Avg. time taken: 7.19476 seconds.
The best 3 of 10 runs took 7.16229 seconds.
2.4.1
Avg. time taken: 7.20860 seconds.
The best 3 of 10 runs took 7.18259 seconds.
We still see big timings for a model that only requires 30
samples. Something is wrong here, and it must be simple.
Initial hypothesis is wrong. Only pushing the front from iteration to iteration has no impact. Somehow the runtime scales with max_samples
. Let's list all the ops that operate on the whole state.sample_collection
:
create_init_state
once per replica._single_thread_ns
1-3 times per parallel replica.
compute_termination
once per major cycle of thread + once at end
count_crossed_edges
-- argsort, bincount (main culprit)compute_evidence_stats
scales with current sample count used, but can build up if number of samples collected grows.Seems like the terminations conditions needs to be made into an accumulant.
Describe the bug In 2.3 I've refactored the algorithm for improvements but there's a loop where I'm pushing samples between iterations. This must be where slow down is happening. I'll profile it, and determine for sure. I don't think the entire state is needed per iteration so it should be fairly trivial to resolve the slow down. Slow down happens around
max_samples>1e6
.JAXNS version
2.3.x