bayesianpymcmcmchierarchical-bayesianarviz

Cannot load a previously saved trace for future inferences using Pymc 3.11.4


I am trying to implement a simple hierarchical bayesian inference. I am inferring a parameter 'm' of a simple linear model y = mx using Pymc. I want to save the trace so that I can load it later. The purpose of loading the trace later is to infer another parameter c that is now part of the new model y = mx+c for a different data set while using the already estimated posterior of m from the previous trace(I am not sure how this is possible, but that is another question). I have tried the following code to save the trace.

import pymc3 as pm
import numpy as np
import random
import arviz as az
import matplotlib.pyplot as plt


np.random.seed(0)

x_data = np.linspace(1, 100, 100)
y_data = 2 * x_data + np.random.uniform(0, 0.9, 100)  

with pm.Model() as linear_model:
    m = pm.Normal('m', mu=0, sigma=10)  
    
    y_pred = m * x_data
    likelihood = pm.Normal('y_obs', mu=y_pred, sigma=1.0, observed=y_data)  
    
    trace = pm.sample(2000, tune=1000)  

    pm.save_trace(trace, 'my_trace_file',overwrite=True)

Saving the trace works. But when I try to load the trace using the following code, I get the error that "TypeError: No model on context stack."

trace = pm.load_trace('my_trace_file')
print("Posterior mean of m:", trace['m'].mean())

Any help or tips would be much helpful as I am new to bayesian inference for parameter estimation and using Pymc3.


Solution

  • The model context is the line with

    with pm.Model() as linear_model:
        ...
    

    Unfortunately, PyMC does not yet provide a simple and easy way to save the model to a file (like arviz does for inference data). However, there is this ModelBuilder. I assume, if you follow this example, you should be able to load your model from a file, and then use it as the context for loading your stored inference data. Something like this stub:

    from pymc_experimental.model_builder import ModelBuilder
    
    class MyModel(ModelBuilder):
        # add more stuff here
        ...
    
    ...
    
    # do stuff and save inference
    
    ...
    

    In another script, you should then be able to load and re-use your inference model/data.

    
    fname = "my_model.nc"
    saved_model = MyModel.load(fname)
    
    # now use the model in this context
    with saved_model:
        trace = pm.load_trace('my_trace_file')