PyMC3 (or Theano) allocating too much virtual memory while sampling

Viewed 186

I'm running a PyMC3 sampler on a remote cluster with a total of 3000 tuning steps, 3000 draws, and 2 chains. The data that I'm trying to fit consists of ~6000 data points and I'm setting up 7 variables (I think) to be optimized by the sampler. Maybe around 20 steps into the sampling, my program crashes because it hits the virtual memory allotment of 15G. I can increase the allotment but then my program sits in the job queue on the cluster for much longer before running, which is not ideal.

I'm relatively unfamiliar with the intricacies of PyMC3 and Theano, but I suspect one of these two is responsible. I've included the code for my PyMC3 model (which incorporates a separate astrophysics program, exoplanet) and the code where I set up the sampling. I have the Theano automatic garbage collection already turned on, but it doesn't seem to help much. Thanks for any help you can provide!

This is part of my current script, utilizing PyMC3 + Theano with another astrophysics program exoplanet:

import numpy as np
import pymc3 as pm
import exoplanet as xo

if mask is None:
    mask = np.ones(len(x), dtype=bool)
with pm.Model() as model:

    mean = pm.Normal("mean", mu=0.0, sd=1.0)

    BoundedNormal_t0 = pm.Bound(pm.Normal, lower=t0_true-0.5, upper=t0_true+0.5)
    t0 = BoundedNormal_t0("t0", mu=t0_true, sd=1.0, shape=1)

    BoundedNormal_per = pm.Bound(pm.Normal, lower=(per_true*0.9), upper=(per_true*1.1))
    period = BoundedNormal_per("period", mu=per_true, sd=1.0, shape=1)

    u = pm.Uniform("u", lower=0.0, upper=1.0, shape=2)
    rho = pm.Uniform("rho", lower=0.0, upper=1000, shape=1)

    BoundedNormal_logr = pm.Bound(pm.Normal, lower=-10.0, upper=0.0)
    logr = BoundedNormal_logr("logr", mu=np.log(rp_rs_true), sd=10.0, shape=1)
    r = pm.Deterministic("r", tt.exp(logr))

    b = xo.distributions.ImpactParameter("b", ror=r, shape=1)

    # Set up a Keplerian orbit for the planets
    orbit = xo.orbits.KeplerianOrbit(period=period, t0=t0, b=b, rho_star=rho)#, r_star=rstar)#, ecc=ecc, omega=omega)

    # Compute the model light curve using starry
    light_curves = xo.LimbDarkLightCurve(u).get_light_curve(
        orbit=orbit, r=r, t=x[mask], texp=texp)
    light_curve = pm.math.sum(light_curves, axis=-1) + mean

    # The likelihood function assuming known Gaussian uncertainty
    pm.Normal("obs", mu=light_curve, sd=yerr[mask], observed=y[mask], dtype='float32')

    map_soln = xo.optimize(start=start)
    map_soln = xo.optimize(start=map_soln, vars=[logr], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[b], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[period, t0], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[u], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[logr], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[b], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[rho], verbose=False)
    map_soln = xo.optimize(start=map_soln, vars=[mean], verbose=False)
    map_soln = xo.optimize(start=map_soln, verbose=False)

np.random.seed(42)
with model:
    tr = pm.sample(chains=2
                   start=map_soln,
                   step=xo.get_dense_nuts_step(target_accept=0.95,start=map_soln),
                   tune=3000,
                   draws=3000)
0 Answers
Related