[WIP] Reduce memory usage in log_likelihood io_pymc3 #1082

nitishp25 · 2020-02-20T16:19:34Z

Description

Related to issue #1077

Checklist

Follows official PR format
Code style correct (follows pylint and black guidelines)
Changes are listed in changelog

arviz/utils.py

arviz/data/io_pymc3.py

OriolAbril

LGTM, It would be great to run some benchmarks and also add some tests of log_likelihood argument behaviour (better wait until #1080 to add tests though, otherwise you will probably have merging issues).

arviz/data/io_pymc3.py

nitishp25 · 2020-02-23T14:21:21Z

These are the benchmarks I ran using memory profiler on this code.

The first 4 graphs show the memory being used with time for samples of different chains.

The last graph summarizes the maximum memory used for different chains and the memory difference between the Preallocated code and Non-preallocated code.

I'll add the tests soon.

Thanks to @OriolAbril for all the help!

Summary

OriolAbril · 2020-02-23T15:54:19Z

This looks great, thanks! It would also be great if you could review @rpgoldman given that it uses _DefaultTrace (only if you have time!).

2 things left to do:

Could you add some tests to check log_likelikhood argument behaves as suposed to? There is no need to check the memory usage in a test. maybe a pytest.mark.parametrize added to test_multiple_observed_rv to check that default returns both observed variables, false returns no log_likelihood group and using a str returns only the selected variable. check_multiple_attrs allows to ensure both cases, either an attr being present or not, the help (and other cases where it is used) should be clear.

On second thought, using pymc3's _DefaultTrace would force users to have a very recent pymc3 version, maybe it is not ideal

nitishp25 · 2020-02-23T16:03:13Z

Could you add some tests to check log_likelikhood argument behaves as suposed to? There is no need to check the memory usage in a test. maybe a pytest.mark.parametrize added to test_multiple_observed_rv to check that default returns both observed variables, false returns no log_likelihood group and using a str returns only the selected variable. check_multiple_attrs allows to ensure both cases, either an attr being present or not, the help (and other cases where it is used) should be clear.

Yes I'm working on testing these cases only. Thanks!

On second thought, using pymc3's defaultTrace would force users to have a very recent pymc3 version, maybe it is not ideal

Then the previous one with the code defined in utils would be more preferable?

rpgoldman · 2020-02-23T16:28:49Z

I'll try to check in the next day or so. Busy today!

rpgoldman · 2020-02-23T16:38:26Z

Minor nit: it's _DefaultTrace that is used, not defaultTrace -- took me a few minutes to figure this out.

OriolAbril · 2020-02-23T18:37:19Z

arviz/tests/test_data_pymc.py

        test_dict = {
            "posterior": ["x"],
            "observed_data": ["y1", "y2"],
            "log_likelihood": ["y1", "y2"],
            "sample_stats": ["diverging", "lp"],
        }
        fails = check_multiple_attrs(test_dict, inference_data)
-        assert not fails
+        if log_likelihood is True:


I'd be better to modify the test_dict, something like:

if not log_likelihood: test_dict.pop("log_likelihood") test_dict["~log_likelihood"] = [] if isinstance(log_likelihood, list): test_dict["log_likelihood"] = ["y1", "~y2"]

OriolAbril · 2020-02-23T18:38:09Z

Also, make sure to get your branch up to date with master!

OriolAbril · 2020-02-23T19:14:44Z

arviz/tests/external_tests/test_data_pymc.py

@@ -215,12 +216,19 @@ def test_multiple_observed_rv(self):
            pm.Normal("y2", x, 1, observed=y2_data)
            trace = pm.sample(100, chains=2)
            inference_data = from_pymc3(trace=trace)
+            inference_data = from_pymc3(trace=trace, log_likelihood=log_likelihood)


It looks like the conversion is done twice.

OriolAbril · 2020-02-23T19:41:02Z

arviz/data/io_pymc3.py

@@ -193,7 +202,7 @@ def sample_stats_to_xarray(self):
    @requires("model")
    def log_likelihood_to_xarray(self):
        """Extract log likelihood and log_p data from PyMC3 trace."""
-        if self.predictions:
+        if self.predictions or not self.log_likelihood:
            return None
        data = self._extract_log_likelihood()


Extra idea. We could put a

try: data = self._extract... except TypeError: warnings.warn("could not compute log likelihood. log_likelihood group will be omitted. Check your model object or set log_likelihood=False") return None

Which if I am correct should fix several issues such as #395 (it has a minimal example to reproduce so it should be easy to check the issue is fixed) or even pymc-devs/pymc#3728

Yes it does fix both the issues!

nitishp25 · 2020-02-23T21:08:57Z

@rpgoldman Can you explain what the 2 peaks near the middle of the graphs above represent?

They seem to occur at the start of allocation of posterior_predictive. And each of those peaks consists more peaks equal to the no. of chains of the samples.
Similar peaks are observed in allocation of log_likelihood but they are relatively smaller (use lesser memory, also only equal to no. of chains, unlike 2 * no. of chains for posterior_predictive).
If you want I can give you the visuals to refer.

rpgoldman · 2020-02-25T20:50:50Z

@rpgoldman Can you explain what the 2 peaks near the middle of the graphs above represent?

So you are asking about the two spikes that come at the start of the process, as opposed to the two near the end. Correct?

Do we know if they are happening in the building of chain_likelihoods, the allocation of the _TraceDict, log_likelilhood_dict, or somewhere else? It seems like that is the most important question to answer. Once we know that, we will know if it is something inside my _TraceDict code (likely the insertmethod), or if it's the list creations in the loop over the chains.

Can you derive plots like the above where the x axis plots position in the code, instead of clock time?

OriolAbril · 2020-02-25T21:11:34Z

Back in the day when I did some experiments on this, I found the pattern to be: first slope corresponds to model creation and trace allocation, plateau is posterior sampling, the gentle slope is posterior predictive sampling and steep slope is conversion to inference data (which in terms of memory is basically retrieving and allocating log_likelihood data).

This would leave the two spikes at the beginning of posterior predictive sampling or at the end of posterior sampling. While writing this I realized it is not difficult to identify which of the two is happening, so I ran a quick example similar to the one above with just posterior sampling, no posterior predictive sampling. I see these same peaks towards the end of sampling.

rpgoldman · 2020-02-25T21:18:08Z

I changed posterior predictive sampling to pre-allocate memory because it used to be building enormous python lists and then translating them into numpy arrays. Now we pre-allocate the arrays and then just fill them. Posterior predictive sampling is still much slower than I would like (I have an open PR for this), but pre-allocation substantially reduces memory usage. That could be what is causing this issue.

I'm a bit confused about why this question is here on the ArviZ development repo, instead of on the `pymc3-devs1 one.

OriolAbril · 2020-02-25T21:34:08Z

I have done a quick check between pymc3 versions, only posterior sampling, no posterior predictive sampling nor conversion to inference data. This is the result for 6 chains:

It is here because we saw this behaviour while checking that PR did indeed reduce memory usage in log_likelihood_to_xarray and we didn't have much idea of what was happening. Now that the issue seems more clear, it is probably better to move to pymc issue or discourse.

If there are no issues with log_likelihood_to_xarray I'll merge

rpgoldman · 2020-02-25T21:37:46Z

I have done a quick check between pymc3 versions, only posterior sampling, no posterior predictive sampling nor conversion to inference data.

That's a relief! It can't be my fault, then, because I didn't change anything in the normal sampling, only posterior predictive. 😌

If there are no issues with log_likelihood_to_xarray I'll merge

Please do!

OriolAbril · 2020-02-25T23:46:38Z

arviz/data/io_pymc3.py

+        log_likelihood_dict = self.pymc3.sampling._DefaultTrace(  # pylint: disable=protected-access
+            len(self.trace.chains)
+        )


we should wrap this in a try except block, so that when pymc3 version is an old one, the error risen explains that pymc3 should be upgraded or arviz downgraded, otherwise users won't know what hit them.

There is no need to edit requirements file though, as pymc3 development is shown there.

OriolAbril · 2020-02-26T14:36:21Z

arviz/data/io_pymc3.py

+            )
+        except AttributeError:
+            raise AttributeError(
+                "Either upgrade PyMC3 to latest version or downgrade ArviZ for log_likelihood."


"Installed version of ArviZ requires PyMC3>=3.8. Please upgrade with `pip install pymc3>=3.8` " "or `conda install -c conda-forge pymc3>=3.8`."

nitishp25 requested a review from OriolAbril February 20, 2020 16:19

OriolAbril reviewed Feb 20, 2020

View reviewed changes

nitishp25 requested a review from OriolAbril February 21, 2020 08:30

nitishp25 changed the title ~~Reduce memory usage in log_likelihood io_pymc3~~ [WIP] Reduce memory usage in log_likelihood io_pymc3 Feb 21, 2020

OriolAbril reviewed Feb 21, 2020

View reviewed changes

arviz/data/io_pymc3.py Outdated Show resolved Hide resolved

OriolAbril reviewed Feb 23, 2020

View reviewed changes

nitishp25 added 6 commits February 24, 2020 00:11

reduce memory usage and add log_likelihood arg

991b917

lint changes

faf7a3d

use pymc3 class

d617674

minor changes

c059f7e

minor update

c3dc4ed

modified test

5f74099

nitishp25 force-pushed the io-pymc3-memory branch from 3377b66 to 5f74099 Compare February 23, 2020 18:58

nitishp25 requested a review from OriolAbril February 23, 2020 19:01

OriolAbril approved these changes Feb 23, 2020

View reviewed changes

OriolAbril reviewed Feb 23, 2020

View reviewed changes

add warning

390632a

OriolAbril approved these changes Feb 25, 2020

View reviewed changes

nitishp25 added 2 commits February 26, 2020 12:59

add import error

62f53ce

correct error

d1b4d12

OriolAbril reviewed Feb 26, 2020

View reviewed changes

nitishp25 and others added 2 commits February 26, 2020 23:20

update warning

f6e00d0

Update CHANGELOG.md

5000d49

OriolAbril merged commit 0a0f18d into arviz-devs:master Feb 26, 2020

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP] Reduce memory usage in log_likelihood io_pymc3 #1082

[WIP] Reduce memory usage in log_likelihood io_pymc3 #1082

nitishp25 commented Feb 20, 2020

OriolAbril left a comment

nitishp25 commented Feb 23, 2020 •

edited

Loading

OriolAbril commented Feb 23, 2020 •

edited

Loading

nitishp25 commented Feb 23, 2020

rpgoldman commented Feb 23, 2020

rpgoldman commented Feb 23, 2020

OriolAbril Feb 23, 2020

OriolAbril commented Feb 23, 2020

OriolAbril Feb 23, 2020

OriolAbril Feb 23, 2020

nitishp25 Feb 23, 2020

nitishp25 commented Feb 23, 2020 •

edited

Loading

rpgoldman commented Feb 25, 2020

OriolAbril commented Feb 25, 2020

rpgoldman commented Feb 25, 2020

OriolAbril commented Feb 25, 2020

rpgoldman commented Feb 25, 2020

OriolAbril Feb 25, 2020

OriolAbril Feb 26, 2020

[WIP] Reduce memory usage in log_likelihood io_pymc3 #1082

[WIP] Reduce memory usage in log_likelihood io_pymc3 #1082

Conversation

nitishp25 commented Feb 20, 2020

Description

Checklist

OriolAbril left a comment

Choose a reason for hiding this comment

nitishp25 commented Feb 23, 2020 • edited Loading

Summary

OriolAbril commented Feb 23, 2020 • edited Loading

nitishp25 commented Feb 23, 2020

rpgoldman commented Feb 23, 2020

rpgoldman commented Feb 23, 2020

OriolAbril Feb 23, 2020

Choose a reason for hiding this comment

OriolAbril commented Feb 23, 2020

OriolAbril Feb 23, 2020

Choose a reason for hiding this comment

OriolAbril Feb 23, 2020

Choose a reason for hiding this comment

nitishp25 Feb 23, 2020

Choose a reason for hiding this comment

nitishp25 commented Feb 23, 2020 • edited Loading

rpgoldman commented Feb 25, 2020

OriolAbril commented Feb 25, 2020

rpgoldman commented Feb 25, 2020

OriolAbril commented Feb 25, 2020

rpgoldman commented Feb 25, 2020

OriolAbril Feb 25, 2020

Choose a reason for hiding this comment

OriolAbril Feb 26, 2020

Choose a reason for hiding this comment

nitishp25 commented Feb 23, 2020 •

edited

Loading

OriolAbril commented Feb 23, 2020 •

edited

Loading

nitishp25 commented Feb 23, 2020 •

edited

Loading