From 334dfe1ff54c94e2582857b78d09be3f739922c0 Mon Sep 17 00:00:00 2001 From: Adrian Seyboldt Date: Wed, 1 Jul 2020 10:02:49 +0200 Subject: [PATCH 1/4] Add time values as sampler stats for NUTS --- pymc3/step_methods/hmc/base_hmc.py | 15 ++++++++++++++- pymc3/step_methods/hmc/hmc.py | 3 +++ pymc3/step_methods/hmc/nuts.py | 10 ++++++++++ 3 files changed, 27 insertions(+), 1 deletion(-) diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index a426f49111..45dcc1ab11 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/pymc3/step_methods/hmc/base_hmc.py @@ -13,6 +13,7 @@ # limitations under the License. from collections import namedtuple +import time import numpy as np import logging @@ -132,6 +133,9 @@ def _hamiltonian_step(self, start, p0, step_size): def astep(self, q0): """Perform a single HMC iteration.""" + perf_start = time.perf_counter_ns() + process_start = time.process_time_ns() + p0 = self.potential.random() start = self.integrator.compute_state(q0, p0) @@ -166,6 +170,9 @@ def astep(self, q0): hmc_step = self._hamiltonian_step(start, p0, step_size) + perf_end = time.perf_counter_ns() + process_end = time.process_time_ns() + self.step_adapt.update(hmc_step.accept_stat, adapt_step) self.potential.update(hmc_step.end.q, hmc_step.end.q_grad, self.tune) if hmc_step.divergence_info: @@ -191,7 +198,13 @@ def astep(self, q0): if not self.tune: self._samples_after_tune += 1 - stats = {"tune": self.tune, "diverging": bool(hmc_step.divergence_info)} + stats = { + "tune": self.tune, + "diverging": bool(hmc_step.divergence_info), + "perf_counter_diff_ns": perf_end - perf_start, + "process_time_diff_ns": process_end - process_start, + "perf_counter_ns": perf_end, + } stats.update(hmc_step.stats) stats.update(self.step_adapt.stats()) diff --git a/pymc3/step_methods/hmc/hmc.py b/pymc3/step_methods/hmc/hmc.py index 3542f1320d..e4a453c0be 100644 --- a/pymc3/step_methods/hmc/hmc.py +++ b/pymc3/step_methods/hmc/hmc.py @@ -48,6 +48,9 @@ class HamiltonianMC(BaseHMC): 'path_length': np.float64, 'accepted': np.bool, 'model_logp': np.float64, + 'process_time_diff_ns': np.int64, + 'perf_counter_diff_ns': np.int64, + 'perf_counter_ns': np.int64, }] def __init__(self, vars=None, path_length=2., max_steps=1024, **kwargs): diff --git a/pymc3/step_methods/hmc/nuts.py b/pymc3/step_methods/hmc/nuts.py index d2409acca2..9e72909a26 100644 --- a/pymc3/step_methods/hmc/nuts.py +++ b/pymc3/step_methods/hmc/nuts.py @@ -72,6 +72,13 @@ class NUTS(BaseHMC): samples, the step size is set to this value. This should converge during tuning. - `model_logp`: The model log-likelihood for this sample. + - `process_time_diff_ns`: The time it took to draw the sample, as defined + by the python standard library `time.process_time_ns`. This counts all + the CPU time, including worker processes in BLAS and OpenMP. + - `perf_counter_diff_ns`: The time it took to draw the sample, as defined + by the python standard library `time.perf_counter_ns` (wall time). + - `perf_counter_ns`: The value of the `time.perf_counter_ns` after drawing + the sample. References ---------- @@ -96,6 +103,9 @@ class NUTS(BaseHMC): "energy": np.float64, "max_energy_error": np.float64, "model_logp": np.float64, + "process_time_diff_ns": np.int64, + "perf_counter_diff_ns": np.int64, + "perf_counter_ns": np.int64, } ] From 07cc7b69e5c962a510633b534eea026f1769a5af Mon Sep 17 00:00:00 2001 From: Adrian Seyboldt Date: Wed, 1 Jul 2020 12:07:18 +0200 Subject: [PATCH 2/4] Use float time counters for nuts stats --- pymc3/step_methods/hmc/base_hmc.py | 14 +++++++------- pymc3/step_methods/hmc/hmc.py | 6 +++--- pymc3/step_methods/hmc/nuts.py | 18 +++++++++--------- pymc3/tests/test_step.py | 12 ++++++++---- 4 files changed, 27 insertions(+), 23 deletions(-) diff --git a/pymc3/step_methods/hmc/base_hmc.py b/pymc3/step_methods/hmc/base_hmc.py index 45dcc1ab11..7bbc722dee 100644 --- a/pymc3/step_methods/hmc/base_hmc.py +++ b/pymc3/step_methods/hmc/base_hmc.py @@ -133,8 +133,8 @@ def _hamiltonian_step(self, start, p0, step_size): def astep(self, q0): """Perform a single HMC iteration.""" - perf_start = time.perf_counter_ns() - process_start = time.process_time_ns() + perf_start = time.perf_counter() + process_start = time.process_time() p0 = self.potential.random() start = self.integrator.compute_state(q0, p0) @@ -170,8 +170,8 @@ def astep(self, q0): hmc_step = self._hamiltonian_step(start, p0, step_size) - perf_end = time.perf_counter_ns() - process_end = time.process_time_ns() + perf_end = time.perf_counter() + process_end = time.process_time() self.step_adapt.update(hmc_step.accept_stat, adapt_step) self.potential.update(hmc_step.end.q, hmc_step.end.q_grad, self.tune) @@ -201,9 +201,9 @@ def astep(self, q0): stats = { "tune": self.tune, "diverging": bool(hmc_step.divergence_info), - "perf_counter_diff_ns": perf_end - perf_start, - "process_time_diff_ns": process_end - process_start, - "perf_counter_ns": perf_end, + "perf_counter_diff": perf_end - perf_start, + "process_time_diff": process_end - process_start, + "perf_counter_start": perf_start, } stats.update(hmc_step.stats) diff --git a/pymc3/step_methods/hmc/hmc.py b/pymc3/step_methods/hmc/hmc.py index e4a453c0be..9c7a533461 100644 --- a/pymc3/step_methods/hmc/hmc.py +++ b/pymc3/step_methods/hmc/hmc.py @@ -48,9 +48,9 @@ class HamiltonianMC(BaseHMC): 'path_length': np.float64, 'accepted': np.bool, 'model_logp': np.float64, - 'process_time_diff_ns': np.int64, - 'perf_counter_diff_ns': np.int64, - 'perf_counter_ns': np.int64, + 'process_time_diff': np.float64, + 'perf_counter_diff': np.float64, + 'perf_counter_start': np.float64, }] def __init__(self, vars=None, path_length=2., max_steps=1024, **kwargs): diff --git a/pymc3/step_methods/hmc/nuts.py b/pymc3/step_methods/hmc/nuts.py index 9e72909a26..30b0705cc8 100644 --- a/pymc3/step_methods/hmc/nuts.py +++ b/pymc3/step_methods/hmc/nuts.py @@ -72,13 +72,13 @@ class NUTS(BaseHMC): samples, the step size is set to this value. This should converge during tuning. - `model_logp`: The model log-likelihood for this sample. - - `process_time_diff_ns`: The time it took to draw the sample, as defined - by the python standard library `time.process_time_ns`. This counts all + - `process_time_diff`: The time it took to draw the sample, as defined + by the python standard library `time.process_time`. This counts all the CPU time, including worker processes in BLAS and OpenMP. - - `perf_counter_diff_ns`: The time it took to draw the sample, as defined - by the python standard library `time.perf_counter_ns` (wall time). - - `perf_counter_ns`: The value of the `time.perf_counter_ns` after drawing - the sample. + - `perf_counter_diff`: The time it took to draw the sample, as defined + by the python standard library `time.perf_counter` (wall time). + - `perf_counter_start`: The value of `time.perf_counter` at the beginning + of the computation of the draw. References ---------- @@ -103,9 +103,9 @@ class NUTS(BaseHMC): "energy": np.float64, "max_energy_error": np.float64, "model_logp": np.float64, - "process_time_diff_ns": np.int64, - "perf_counter_diff_ns": np.int64, - "perf_counter_ns": np.int64, + "process_time_diff": np.float64, + "perf_counter_diff": np.float64, + "perf_counter_start": np.float64, } ] diff --git a/pymc3/tests/test_step.py b/pymc3/tests/test_step.py index 2e4ce427a2..28448b8cad 100644 --- a/pymc3/tests/test_step.py +++ b/pymc3/tests/test_step.py @@ -979,7 +979,7 @@ def test_linalg(self, caplog): def test_sampler_stats(self): with Model() as model: - x = Normal("x", mu=0, sigma=1) + Normal("x", mu=0, sigma=1) trace = sample(draws=10, tune=1, chains=1) # Assert stats exist and have the correct shape. @@ -995,6 +995,9 @@ def test_sampler_stats(self): "step_size_bar", "tree_size", "tune", + "perf_counter_diff", + "perf_counter_start", + "process_time_diff", } assert trace.stat_names == expected_stat_names for varname in trace.stat_names: @@ -1002,7 +1005,8 @@ def test_sampler_stats(self): # Assert model logp is computed correctly: computing post-sampling # and tracking while sampling should give same results. - model_logp_ = np.array( - [model.logp(trace.point(i, chain=c)) for c in trace.chains for i in range(len(trace))] - ) + model_logp_ = np.array([ + model.logp(trace.point(i, chain=c)) + for c in trace.chains for i in range(len(trace)) + ]) assert (trace.model_logp == model_logp_).all() From fa87be3fa3acd40800a299c73596b503bb904870 Mon Sep 17 00:00:00 2001 From: Adrian Seyboldt Date: Wed, 1 Jul 2020 15:54:47 +0200 Subject: [PATCH 3/4] Add timing sampler stats to release notes --- RELEASE-NOTES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 30f96ada6c..ac3282f5a3 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -3,6 +3,7 @@ ## PyMC3 3.9.x (on deck) ### Documentation - Notebook on [multilevel modeling](https://docs.pymc.io/notebooks/multilevel_modeling.html) has been rewritten to showcase ArviZ and xarray usage for inference result analysis (see [#3963](https://github.com/pymc-devs/pymc3/pull/3963)) +- Add sampler stats `process_time_diff`, `perf_counter_diff` and `perf_counter`, that record wall and CPU time for each NUTS and HMC sample (see [ #3986](https://github.com/pymc-devs/pymc3/pull/3986)). ## PyMC3 3.9.2 (24 June 2020) ### Maintenance From dc892022b801b0bf38800a6e9508a6d059a40cc1 Mon Sep 17 00:00:00 2001 From: Adrian Seyboldt Date: Wed, 1 Jul 2020 17:55:26 +0200 Subject: [PATCH 4/4] Improve doc of time related sampler stats Co-authored-by: Alexandre ANDORRA --- RELEASE-NOTES.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index ac3282f5a3..c678272efb 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -3,7 +3,9 @@ ## PyMC3 3.9.x (on deck) ### Documentation - Notebook on [multilevel modeling](https://docs.pymc.io/notebooks/multilevel_modeling.html) has been rewritten to showcase ArviZ and xarray usage for inference result analysis (see [#3963](https://github.com/pymc-devs/pymc3/pull/3963)) -- Add sampler stats `process_time_diff`, `perf_counter_diff` and `perf_counter`, that record wall and CPU time for each NUTS and HMC sample (see [ #3986](https://github.com/pymc-devs/pymc3/pull/3986)). + +### New features +- Add sampler stats `process_time_diff`, `perf_counter_diff` and `perf_counter_start`, that record wall and CPU times for each NUTS and HMC sample (see [ #3986](https://github.com/pymc-devs/pymc3/pull/3986)). ## PyMC3 3.9.2 (24 June 2020) ### Maintenance