Very slow NUTS HMC

I am a new user of Pyro. I wrote the below script, however, it takes too long to run. The example dataset in the script is very small, and it takes ~65 seconds to run it. I am planning to feed 500 times larger dataset. I am sure the way I constructed this model is sufficient, but I do not know how to overcome it.

import pdb
import pyro
from pyro.distributions import Beta, Exponential, Normal
import torch
import time
start_time = time.time()

DATA = [
    # user, item, order, rating
    [0, 0, 0, 2.5],
    [0, 1, 1, 4.0],
    [0, 2, 2, 5.0],
    [1, 0, 0, 5.0],
    [1, 1, 1, 5.0],
    [1, 2, 2, 5.0],
    [2, 0, 0, 7.5],
    [2, 1, 1, 5.0],
    [2, 2, 2, 2.5],
    [3, 0, 0, 7.5],
    [3, 1, 1, 5.0],
    [3, 2, 2, 2.5],
    [4, 0, 0, 7.5],
    [4, 1, 1, 5.0],
    [4, 2, 2, 2.5],
    [5, 0, 0, 7.5],
    [5, 1, 1, 5.0],
    [5, 2, 2, 2.5],
    [6, 0, 0, 7.5],
    [6, 1, 1, 5.0],
    [6, 2, 2, 2.5],
    [7, 0, 0, 7.5],
    [7, 1, 1, 5.0],
    [7, 2, 2, 2.5]
]


K = len(list(set([int(row[0]) for row in DATA])))
N = len(list(set([int(row[1]) for row in DATA])))
R = len(DATA)
print((K, N, R))

def model():
    # Sample a cultural consensus for each item
    consensus = []
    for i in pyro.plate("consensus", N):
        consensus.append(pyro.sample("consensus_{}".format(i), Beta(1, 1)))

    # Sample a cultural competence for each user
    competence = []
    for i in pyro.plate("competence", K):
        competence.append(pyro.sample("competence_{}".format(i), Exponential(0.25)))

    # Sample a learning rate for each user
    learningRate = []
    for i in pyro.plate("learning_rate", K):
        learningRate.append(pyro.sample("learning_rate{}".format(i), Exponential(0.5)))
    # Sample each rating
    rating = []
    for i in pyro.plate("data_loop", R):
        alpha = consensus[int(DATA[i][1])] * competence[int(DATA[i][0])] + learningRate[int(DATA[i][0])] * int(DATA[i][2])
        beta = (1 - consensus[int(DATA[i][1])]) * competence[int(DATA[i][0])] + learningRate[int(DATA[i][0])] * int(DATA[i][2])
        rating.append(pyro.sample("rating_{}".format(i), Beta(alpha, beta)))



dataDict = {"rating_{}".format(i):torch.tensor(DATA[i][3]/10.0) for i in range(R)}
conditionedModel = pyro.condition(model, data=dataDict)
nutsKernel = pyro.infer.NUTS(conditionedModel, adapt_step_size=True)
mcmc = pyro.infer.MCMC(nutsKernel, num_samples=100, warmup_steps=50, num_chains=31)
mcmc.run()

samples = mcmc.get_samples()
mcmc.diagnostics()
mcmc.summary()
print ("The script took", time.time() - start_time, "to run")Preformatted text
  • the for loops you are using get executed sequentially and are thus quite slow. for faster code you need to vectorize your code and use a pyro.plate as described e.g. here
  • the nuts implementation in numpyro can be quite a bit faster than that in pyro, especially for smaller models, so you might consider that if you’re especially concerned with speed
2 Likes