NUTS sampler needs more than 6 GB RAM on model with just 26 unkowns

[bjornsing]

I'm trying to build a relatively simple model in pymc3 (version 3.0rc1) but my script (see below) hangs when instantiating the NUTS sampler. Top tells me it's because the python process needs more than 5.9 GB RAM, starts paging to disk and grinds to a halt. To me this seems like an extraordinary amount of memory given that the model contains only 26 unknowns.

The code:

import numpy as np
import theano
assert(theano.config.device == 'cpu'), "This code should probably run on the cpu..."
assert(theano.config.floatX == 'float64'), "This code need double precision floats..."
from pymc3 import Model, Uniform, Normal, Gamma, StudentT

# Length of core in cm:
L = 100

# Radiocarbon dates (depth in cm and age in years):
RCDs = [{'depth': 45.0, 'age': 3200.0, 'sd': 130},
        {'depth': 60.0, 'age': 5700.0, 'sd': 130},
        {'depth': 90.0, 'age': 7500.0, 'sd': 130},]

# Number of segments:
K = 25

# Year core was retrieved:
A = 2016

# MODEL SETUP - Essentially and attempt to implement https://projecteuclid.org/euclid.ba/1339616472

# Standard deviation in top layer age:
sigma_A = 10

# Gamma parameters:
# TBD: Shouldn't these parameters be "learnt" from the data?
a_alpha = 2
b_alpha = 2.0/1.1

bacon_model = Model()

with bacon_model:

    # TBD: Shouldn't this be a truncated normal distribution?
    theta = Normal('theta', mu=A, sd=sigma_A)

    # TODO: It seems prior on w shouldn't be uniform, but what should it be...?
    w = Uniform('w', lower=0.0, upper=1.0)

    alpha = [Gamma('a_' + str(j), alpha=a_alpha, beta=b_alpha) for j in range(K)]

    def x(j):
        if j <= K:
            return w * x(j+1) + (1.0 - w) * alpha[j - 1]
        else:
            # TBD: Is this correct...?
            return 0

    delta_c = float(L)/float(K)

    def t(z):
        depth = 0
        age   = theta

        j = 1
        while depth < (z - delta_c):
            depth += delta_c
            age   += x(j) * delta_c
            j     += 1

        age += x(j) * (z - depth)

        return age

    for RCD in RCDs:
        # TODO: Switch from Normal to StudentT...
        RCD['variable'] = Normal('RCD_' + str(RCD['depth']), mu=t(RCD['depth']), sd=RCD['sd'], observed=RCD['age'])

# INFERENCE - Copy-pasted from https://peerj.com/preprints/1686.pdf

from scipy import optimize
from pymc3 import find_MAP, NUTS, sample
import datetime

with bacon_model:

    print "*** Finding MAP (%s) ***" % datetime.datetime.now()

    start = find_MAP(fmin=optimize.fmin_powell)

    print "MAP:", start

    print "*** Instantiating NUTS sampler (%s) ***" % datetime.datetime.now()

    step = NUTS(scaling=start)

    # NOT REACHED,
    # because NUTS._init_ seems to need more than 5.9 GB RAM and system starts swapping...
    #
    # Relevant output line from top:
    #   PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND
    #26290 xxxxx     20   0 10.233g 5.901g   2600 D  12.6 76.8  48:50.91 python

    print "*** Ready to sample (%s) ***" % datetime.datetime.now()

    trace = sample(2000, step, start=start)

    print "*** Done sampling (%s) ***" % datetime.datetime.now()

Since I've never managed to run the model it may not be correct, but in any case I find it rather strange that it needs so much memory.

The version of pymc3 I'm using:

$ pip show pymc3

---
Metadata-Version: 2.0
Name: pymc3
Version: 3.0rc1
Summary: PyMC3
Home-page: http://github.com/pymc-devs/pymc
Author: Thomas Wiecki
Author-email: thomas.wiecki@gmail.com
Installer: pip
License: Apache License, Version 2.0
Location: /home/coach/local/miniconda2/lib/python2.7/site-packages
Requires: nbsphinx, CommonMark, joblib, numpy, patsy, sphinx, pandas, scipy, numpydoc, matplotlib, theano, recommonmark
Classifiers:
  Development Status :: 5 - Production/Stable
  Programming Language :: Python
  Programming Language :: Python :: 2
  Programming Language :: Python :: 3
  Programming Language :: Python :: 2.7
  Programming Language :: Python :: 3.4
  Programming Language :: Python :: 3.5
  License :: OSI Approved :: Apache Software License
  Intended Audience :: Science/Research
  Topic :: Scientific/Engineering
  Topic :: Scientific/Engineering :: Mathematics
  Operating System :: OS Independent

[ColCarroll]

Thanks for the awesome reproducible example!

I'm trying to run it locally -- I suspect that the memory usage might be due to the recursive implementation that recently became iterative.

Running this on my machine (I'm on python3.5, not 2.7) actually ends up hanging on find_MAP, and eventually dumping a super cryptic and long log message (rerunning now to get a log I can post).

If it is easy, how long does the find_MAP step take for you?

[bjornsing]

You're more than welcome!

find_MAP takes about 20 min on my machine (intel core i7 something).

[bjornsing]

I just tried pip install git+https://github.com/pymc-devs/pymc3, but it behaves similarly: finishes find_MAP in about 25 min with approximately 1 GB RAM, but then it goes into NUTS._init_ and slowly grows to about 6 GB RES (and 15+ GB VIRT). Since I've only got 8 GB RAM my system starts swapping to/from disk and python grinds to a halt:

  PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND
32300 bjornsm+  20   0 18.280g 6.330g   1044 D   3.7 82.3  52:22.55 python

[bjornsing]

I did a pip install --upgrade git+https://github.com/pymc-devs/pymc3 and then left the script above running over night. This morning top had this to say:

  PID USER      PR  NI    VIRT    RES    SHR S  %CPU %MEM     TIME+ COMMAND                                                                                        
24988 bjornsm+  20   0 24.371g 6.369g    748 D   3.0 82.8 114:14.47 python

Unfortunately I don't know exactly how far it got because I redirected standard output to a file, and it seems that the output was buffered.

[fonnesbeck]

Can you check to see how many Theano nodes are in the model's graph? If its creating a lot of nodes, this will take up a lot of memory even without data.

[bjornsing]

How can I check that...? I tried googling for "pymc3 theano number of nodes" but couldn't find anything relevant.

[ColCarroll]

The question is about memory usage rather than timing, so this is sort of beside the point, I modified @bjornsing 's script to time and log the four steps above for different values of K.

Raw data is at https://gist.github.com/ColCarroll/307a7598e9063dbb58dbeee1d32e916f

[twiecki]

I bet that's estimating the scaling. Can you try passing in a covariance matrix (e.g. identity, or via ADVI) directly?

[bjornsing]

Thanks! I tried passing a covariance vector(?) derived with ADVI to the NUTS initializer (replacing everything under INFERENCE in my post above):

# INFERENCE

from scipy import optimize
from pymc3 import variational, NUTS, sample
import datetime

with bacon_model:

    print "*** Finding variational ELBO (%s) ***" % datetime.datetime.now()

    mu, sds, elbo = variational.advi(n=70000)

    print "*** Initializing NUTS (%s) ***" % datetime.datetime.now()

    step = NUTS(scaling=bacon_model.dict_to_array(sds), is_cov=True)

    print "*** Ready to sample (%s) ***" % datetime.datetime.now()

    trace = sample(2000, step, start=mu, progressbar=True)

    print "*** Done sampling (%s) ***" % datetime.datetime.now()

But I also bought 16 GB more RAM and upgraded to the latest git snapshot, so I'm not entirely certain that it was the changes above that got me this result:

*** Finding variational ELBO (2016-10-13 21:27:28.626357) ***
Iteration 0 [0%]: ELBO = -1290.12
Iteration 7000 [10%]: Average ELBO = -1242.57
Iteration 14000 [20%]: Average ELBO = -1219.94
Iteration 21000 [30%]: Average ELBO = -1186.98
Iteration 28000 [40%]: Average ELBO = -1172.12
Iteration 35000 [50%]: Average ELBO = -1166.98
Iteration 42000 [60%]: Average ELBO = -1164.93
Iteration 49000 [70%]: Average ELBO = -1163.8
Iteration 56000 [80%]: Average ELBO = -1163.44
Iteration 63000 [90%]: Average ELBO = -1162.76
Finished [100%]: Average ELBO = -1162.17
*** Initializing NUTS (2016-10-13 21:38:55.584747) ***
*** Ready to sample (2016-10-13 22:04:02.184681) ***
100%|███████████████████████████████████████| 2000/2000 [01:27<00:00, 25.25it/s]
*** Done sampling (2016-10-13 22:05:30.028320) ***

In any case, works for me. :P But I hear (Py)Stan can do this in a matter of minutes...

[twiecki]

It should be bacon_model.dict_to_array(sds)**2 to convert to a covariance matrix.

Seems like sampling only takes less than 2 minutes?

[bjornsing]

Thanks! I changed the initialization to NUTS(scaling=bacon_model.dict_to_array(sds)**2, is_cov=True) and got this result:

*** Finding variational ELBO (2016-10-14 16:25:45.661955) ***
Iteration 0 [0%]: ELBO = -1278.42
Iteration 7000 [10%]: Average ELBO = -1240.36
Iteration 14000 [20%]: Average ELBO = -1223.07
Iteration 21000 [30%]: Average ELBO = -1187.61
Iteration 28000 [40%]: Average ELBO = -1171.84
Iteration 35000 [50%]: Average ELBO = -1166.92
Iteration 42000 [60%]: Average ELBO = -1165.17
Iteration 49000 [70%]: Average ELBO = -1164.18
Iteration 56000 [80%]: Average ELBO = -1163.27
Iteration 63000 [90%]: Average ELBO = -1163.13
Finished [100%]: Average ELBO = -1162.36
*** Initializing NUTS (2016-10-14 16:37:11.536609) ***
*** Ready to sample (2016-10-14 17:01:34.673211) ***
100%|███████████████████████████████████████| 2000/2000 [01:16<00:00, 26.30it/s]
*** Done sampling (2016-10-14 17:02:50.726031) ***

And yes, sampling seems super quick (around a minute), while finding the ELBO and initializing NUTS takes about 35 min. I can understand it takes around 10 min to find the ELBO with variational.advi(n=70000), but why does it take over 20 min to initialize NUTS with a "known" scaling...?

(The reason I ask is because I plan to built a much more complex model and rent a machine with a lot of cores to run it on, but if 80% of the time is in initialization it doesn't make sense to try to parallelize...)

[twiecki]

@bjornsing The main problem I see is that your model uses loops which will incur a much greater memory cost. Rather, try to vectorize your model using the shape kwarg.

[bjornsing]

@twiecki Any pointers / examples on how to vectorize my model using the shape kwarg...? I tried googling but can't find much...

[twiecki]

Replace your loops like this

alpha = [Gamma('a_' + str(j), alpha=a_alpha, beta=b_alpha) for j in range(K)]
->
alpha = Gamma('a', alpha=a_alpha, beta=b_alpha, shape=K)