Saving and loading models

Saving and loading models¶

In this short how-to, tutorial, we show how to save a HSSM model instance and its inference results to disk and then re-instantiate the model from the saved files.

Load data and instantiate HSSM model¶

In [1]:

Copied!





import hssm

cav_data = hssm.load_data("cavanagh_theta")

basic_hssm_model = hssm.HSSM(
    data=cav_data,
    process_initvals=True,
    link_settings="log_logit",
    model="angle",
    include=[
        {
            "name": "v",
            "formula": "v ~ 1 + C(stim)",
        }
    ],
)
import hssm

cav_data = hssm.load_data("cavanagh_theta")

basic_hssm_model = hssm.HSSM(
    data=cav_data,
    process_initvals=True,
    link_settings="log_logit",
    model="angle",
    include=[
        {
            "name": "v",
            "formula": "v ~ 1 + C(stim)",
        }
    ],
)

Model initialized successfully.

Generate inference results¶

MCMC¶

In [2]:

Copied!





basic_hssm_model.sample(sampler="nuts_numpyro",
                        tune=100,
                        draws=100,
                        chains=2)
basic_hssm_model.sample(sampler="nuts_numpyro",
                        tune=100,
                        draws=100,
                        chains=2)

Using default initvals.

/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.11/site-packages/pymc/sampling/jax.py:475: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(
sample: 100%|██████████| 200/200 [01:06<00:00,  3.01it/s, 127 steps of size 3.16e-02. acc. prob=0.96]
sample: 100%|██████████| 200/200 [00:57<00:00,  3.47it/s, 127 steps of size 2.55e-02. acc. prob=0.92]
We recommend running at least 4 chains for robust computation of convergence diagnostics
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
100%|██████████| 200/200 [00:00<00:00, 306.78it/s]

Out[2]:

arviz.InferenceData

posterior

log_likelihood

sample_stats

observed_data

VI¶

Saving and Loading the model¶

In [4]:

Copied!

basic_hssm_model.save_model(model_name="test_model")
basic_hssm_model.save_model(model_name="test_model")

We are using the defaults here, which save the model and its inference results to the hssm_models/test_model/ directory inside your curerent working directory.

Up to three files are saved in the model directory:

model.pkl: The model instance.
traces.nc: The MCMC traces.
vi_traces.nc: The VI traces.

We can now load the model from the directory we just created, using the HSSM classmethod load_model.

In [5]:

Copied!

loaded_model = hssm.HSSM.load_model(path="hssm_models/test_model")
loaded_model = hssm.HSSM.load_model(path="hssm_models/test_model")

Model initialized successfully.

With this simple workflow your models are portable across sessions and machines.