NUTS is slow

[alxempirical]

I'm running the script below with 2,000 samples. It's surprisingly slow, a bit under 9s / sample under Python 2.7.11 :: Anaconda custom (x86_64) on a Mid-2015 MacBook Pro with 2.5GHz CPUs (four of them, but the sampler seems to be single-threaded.) I'd be grateful for any suggestions about how to speed it up.

from pymc3 import Model, Normal, Gamma, Lognormal, Exponential, NUTS, sample, find_MAP, traceplot

from load_table import walltime, numsims, numconstraints, numtargets


def LogNormalInvGamma(mu, alpha, beta, n, name):
    tau = Gamma('tau_' + name, alpha=alpha, beta=beta)
    mu = Normal('mu_' + name, mu=mu, tau=(n * tau))
    return Lognormal(name, mu=mu, tau=tau)


model = Model()

with model:

    # Equation (3) in the doc string.  But remember, comments go stale!
    C_1 = LogNormalInvGamma(mu=0, alpha=1, beta=1,
                            n=Exponential('C_1_DOF', 1),
                            name='constraints')

    # Equation (4)
    C_2 = LogNormalInvGamma(mu=0, alpha=1, beta=1,
                            n=Exponential('C_2_DOF', 1),
                            name='targets')

    # Now we model C_0 as the Lognormal which varies around the rest of
    # Equation (1) / Equation (5)
    C_0_mu = (C_1 * numconstraints) + (C_2 * numsims * numtargets)
    walltime_model = Lognormal('walltime',
                               mu=C_0_mu,
                               tau=Gamma('tau_const', 1, 1),  # 𝞽₀ in Equ (2)
                               observed=walltime)

    # The section below is run in a separate module, which imports the model constructed above.
    # Probably not relevant, but I don't really know why it's taking so long.

    # Rough computation of MAP for the model.
    map_estimate = find_MAP(model=model.model, fmin=optimize.fmin_powell)

    # Draw 2,000 posterior samples
    trace = sample(2000, start=map_estimate)

[twiecki]

Often that's due to bad initialization. Can you try to init with ADVI (e.g. http://pymc-devs.github.io/pymc3/notebooks/stochastic_volatility.html#Fit-Model)?

[twiecki]

Oh, and hey @alxempirical, good to see you here and glad you're giving pymc3 a test-drive :).

[twiecki]

For this model Metropolis might also be good enough and probably way faster.

[alxempirical]

good to see you here

You too!

I don't think the problem is bad initialization. Here are the ADVIFit.means values vs the find_MAP values:

                           ADVI        MAP
mu_constraints             -1.4       -1.4
mu_targets                 -6.8       -9.4
C_1_DOF                    0.57       0.84
tau_constraints            0.44       0.81
constraints               0.079      0.079
C_2_DOF                    0.38       0.12
tau_targets               0.053       0.27
targets                 2.5e-05    2.5e-05
tau_const                    14         14

These values fit well within the sample distributions:

Metropolis might also be good enough and probably way faster.

Thanks, I will give that a go.

Best regards,
Alex

[twiecki]

The cov matrix found by ADVI might still be better suited but I agree, find_MAP seems to find a reasonable starting point. How many data points do you pass in?

[alxempirical]

I'm training on 2,000 data points.

How do I access the cov(ariance?) matrix from advi? It's not mentioned as a return value in variational.advi, only means, stds and elbo_vals.

Also, any tips/docs on profiling? I would guess Theano confuses standard python profiling tools, but could be wrong.

[twiecki]

step = pm.NUTS(scaling=model.dict_to_array(sds), is_cov=True) uses the standard deviations as the diagonal. Writing that I realize that it should obviously use the variances so I'll fix that.

There are some profiling tools in pymc3 that @jsalvatier wrote a while back but I haven't tried them. Would need to take a closer look.

is it faster with metropolis or slice?

[fonnesbeck]

See the model.profile method for built-in profiling.

[alxempirical]

Thanks for the suggestions, both of you. I will try them both out when I productionize this code. I've got a decent run of results now, so I'm going to move on to the next step, for a while.

[twiecki]

#1523 and #1522 should address this.

[fonnesbeck]

This still appears to be a problem. Even with ADVI initialization, this is what happens for a relatively simple model:

Ironically, this model used to work a few weeks ago. The model structure is as follows:

   with Model() as model:

        cases, n = inpatient_data[['cases', 'n']].values.T

        # Cases weighted by monitoring effort and market share
        weights = monitoring_rate[INPATIENT] * np.array(inpatient_market_share)[inpatient_data.year]

        # Uniform prior on weighted cases
        weighted_cases = Uniform('weighted_cases', cases, n, shape=inpatient_data.shape[0],
                                         testval=(cases/weights).astype(int))
        
        # Likelihood of observed cases
        weighting = Binomial('weighting', n=weighted_cases.astype('int32'), p=weights, observed=cases)

        if pool_years:
            
            assert not len(cases) % 3
            
            pooled_shape = int(len(cases)/3)
                                    
            pooled_cases = Deterministic('pooled_cases', tt.reshape(weighted_cases, (pooled_shape, 3)).sum(-1))
            
            θ = Normal('θ', 0, sd=1000, shape=pooled_shape)

            λ = Deterministic('λ', tt.exp(θ))

            AGE_obs = Potential('AGE_obs', Poisson.dist(λ*(np.reshape(n, (pooled_shape, 3))[:, 0])).logp(pooled_cases))
                                
        else:
                      
            θ = Normal('θ', 0, sd=1000, shape=inpatient_data.shape[0])

            # Poisson rate parameter
            λ = Deterministic('λ', tt.exp(θ))
                      
            # Poisson likelihood
            AGE_obs = Potential('AGE_obs', Poisson.dist(λ*n).logp(weighted_cases))

[twiecki]

@fonnesbeck maybe do a git bisect to see when this started happening?