How can I create a class using pm.Model's class like keras.Model or torch.nn.Module option?

Questions v5
, ,
2

The API for PyMC is a bit different from Torch. You don’t ever need to subclass pm.Model to create a model. Instead, it is used as a context manager, as follows:

basic_model = pm.Model()

with basic_model:

    # Priors for unknown model parameters
    alpha = pm.Normal("alpha", mu=0, sigma=10)
    beta = pm.Normal("beta", mu=0, sigma=10, shape=2)
    sigma = pm.HalfNormal("sigma", sigma=1)

    # Expected value of outcome
    mu = alpha + beta[0] * X1 + beta[1] * X2

    # Likelihood (sampling distribution) of observations
    Y_obs = pm.Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)

This example comes from the introduction and overview example notebook in the examples gallery, which I highly recommend checking out.

If you wanted to write a class wrapper around a PyMC model, perhaps to get some additional functionality (like a .fit method), you wouldn’t subclass pm.Model. Instead just write a generic class with a pm.Model as a member variable and proceed as normal.

3 Likes