search

google / bayesnf

Bayesian Neural Field models for prediction in large-scale spatiotemporal datasets
https://google.github.io/bayesnf/
Apache License 2.0
23 stars 3 forks source link

Unable to pickle BayesianNeuralFieldMLE #48

Open f0lie opened 2 months ago

f0lie commented 2 months ago 
with open('/content/drive/MyDrive/Colab Notebooks/models.pkl', 'wb') as f:
  pickle.dump(models["states"]["birth_rate"], f)

Results the error message.

AttributeError                            Traceback (most recent call last)
[<ipython-input-87-db6752b3d600>](https://localhost:8080/#) in <cell line: 1>()
      1 with open('/content/drive/MyDrive/Colab Notebooks/models.pkl', 'wb') as f:
----> 2   pickle.dump(models["states"]["birth_rate"], f)

AttributeError: Can't pickle local object 'structtuple.<locals>.StructTuple'
fsaad commented 2 months ago

Thanks for the report, this issue is being tracked here: https://github.com/google/bayesnf/issues/36

ColCarroll commented 2 months ago

Specifically, this is coming from using a JointDistributionCoroutine, which dynamically generates a StructTuple to represent the state, which means pickle can't represent it.

Depending on what kind of object models['states']['birth_rate'] is, you may be able to use arviz

import arviz as az

idata = az.from_pytree(models['states']['birth_rate'])
idata.to_netcdf("models.nc")

models = az.from_netcdf("models.nc")  # To reload