Best practice for parallel model selection, especially avoidance of recompilation

Questions v5
7

Hey @lucidfrontier45,

Since the compiled model can’t be serialized you could get around that by creating multiple Ray actors and only compiling those one time and then in parallel you can swap out the data.

Running on a local Ray cluster would look like this:

import nutpie
import numpy as np
import pymc as pm
import pytensor.tensor as pt
import pytensor
import ray

print(pm.__version__)

# Start a local ray cluster
ray.init()

np.random.seed(0)

w = [1.0, -2.0, 3.0]
b = 4.0
N = 100

X = np.random.randn(N, len(w))
z = np.dot(X, w) + b
y = np.random.poisson(np.exp(z))

# Define your ray actor
@ray.remote
class NutpieModel:
    def __init__(self, X, y):
        with pm.Model() as self.model:
            # data
            X_ = pm.Data("X", X)
            y_ = pm.Data("y", y)

            # Prior: P(w|α)
            w_ = pm.Normal("w", mu=0.0, sigma=5.0, shape=(len(w), ))
            b_ = pm.Normal("b", mu=0.0, sigma=5.0)

            alpha = pm.HalfNormal("alpha", sigma=5.0)

            # Predictor: z = f(x, w)
            z = pt.dot(X_, w_) + b_

            # Likelihood: P(y|z)
            pm.NegativeBinomial("y_obs", mu=pt.exp(z), alpha=alpha, observed=y_)
        self.compiled_model = nutpie.compile_pymc_model(self.model)

    def sample(self, X = None, y = None):
        if X is None and y is None:
            idata = nutpie.sample(self.compiled_model)
        else:
            idata = nutpie.sample(self.compiled_model.with_data({"X": X, "y": y}))
        return idata

# Create 4 Ray actors
nutpie_models = [NutpieModel.remote(X=X, y=y) for _ in range(4)]

# Each actor which will be compiled once is called to run in parallel 10 times
idata_references = [nutpie_model.sample.remote() for nutpie_model in nutpie_models for _ in range(10)]

# Get back the inference data objects
idatas = ray.get(idata_references)

This is just an example and you will need to tailor this to your specific use case but I think this will help you run in parallel while minimizing compilation.

I really hope this helps!