How to correctly load a trained ContinuousApproximator model to sample from it?

General
1

Hi!

I am getting started to work with BayesFlow on the dev branch (latest commit 8210610). Sorry for all the questions! :smiley:
I am trying to load a saved approximator model via keras.saving.load_model("approximator.keras") and sample from it. However, I encounter an error when calling .sample()
Here’s an example based on the linear regression notebook:

import os

if "KERAS_BACKEND" not in os.environ:
	# set this to "torch", "tensorflow", or "jax"
	os.environ["KERAS_BACKEND"] = "torch"

import bayesflow as bf
import keras
import numpy as np

def prior():
	# beta: regression coefficients (intercept, slope)
	beta = np.random.normal([2, 0], [3, 1])
	return dict(beta=beta)

def likelihood(beta):
	# x: predictor variable
	x = np.random.normal(0, 1, size=10)
	# y: response variable
	y = np.random.normal(beta[0] + beta[1] * x, size=10)
	return dict(y=y, x=x)

simulator = bf.simulators.make_simulator([prior, likelihood])
adapter = (
	bf.Adapter()
	.as_set(["x", "y"])
	.standardize()
	.concatenate(["beta"], into="inference_variables")
	.concatenate(["x", "y"], into="summary_variables")
)
inference_network = bf.networks.FlowMatching()
summary_network = bf.networks.DeepSet(depth=2)
approximator = bf.ContinuousApproximator(
   inference_network=inference_network,
   summary_network=summary_network,
   adapter=adapter,
)
epochs = 1
num_batches = 1
batch_size = 8
optimizer = keras.optimizers.Adam(learning_rate=5e-4, clipnorm=1.0)
approximator.compile(optimizer=optimizer)
history = approximator.fit(
	epochs=epochs,
	num_batches=num_batches,
	batch_size=batch_size,
	simulator=simulator,
)

# Save and reload the model
approximator.save("approximator.keras")
loaded_approximator = keras.saving.load_model("approximator.keras")

# Try sampling
num_samples = 1000
val_sims = simulator.sample(200)
conditions = {k: v for k, v in val_sims.items() if k != "beta"}
pdraws = loaded_approximator.sample(conditions=conditions, num_samples=num_samples)

But I get this error message:

---------------------------------------------------------------------------
ValueError                            	Traceback (most recent call last)
File /data/homes/reiser/projects/bayesflow/issue.py:63
 	60 conditions = {k: v for k, v in val_sims.items() if k != "beta"}
 	62 # obtain num_samples samples of the parameter posterior for every valididation dataset
---> 63 pdraws = loaded_approximator.sample(conditions=conditions, num_samples=num_samples)

File ~/projects/bayesflow/bayesflow/approximators/continuous_approximator.py:142, in ContinuousApproximator.sample(self, num_samples, conditions, split, **kwargs)
	134 def sample(
	135 	self,
	136 	*,
   (...)
	140 	**kwargs,
	141 ) -> dict[str, np.ndarray]:
--> 142 	conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
	143 	conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
	144 	conditions = {"inference_variables": self._sample(num_samples=num_samples, **conditions, **kwargs)}

File ~/projects/bayesflow/bayesflow/adapters/adapter.py:76, in Adapter.__call__(self, data, inverse, **kwargs)
 	73 if inverse:
 	74 	return self.inverse(data, **kwargs)
---> 76 return self.forward(data, **kwargs)

File ~/projects/bayesflow/bayesflow/adapters/adapter.py:60, in Adapter.forward(self, data, **kwargs)
 	57 data = data.copy()
...
---> 57     	raise ValueError("Cannot call `forward` with `strict=False` before calling `forward` with `strict=True`.")
 	59 	# copy to avoid side effects
 	60 	data = data.copy()

ValueError: Cannot call `forward` with `strict=False` before calling `forward` with `strict=True`.

It seems like the adapter in loaded_approximator is not properly initialized after loading the model.
What would be the correct way to load the trained approximator model?

Thank you!

2

No need to be sorry, in the contrary, thanks for using (and thereby testing) the dev branch. Your code is correct, you encountered another bug in the (de)serialization, as the concatenate transform accumulates state which is not yet serialized, and therefore missing after deserialization. We will try to add more thorough testing for the serialization of the adapter to better detect those problems in the future, and supply a fix for the concatenate transform.

3

I have merged a fix to the transform, and somewhat more stringent tests for the rest of the transforms are on their way.
Please let us know if this fixes your problem, and if you encounter anything else, please reach out again, this kind of feedback is really valuable!

4

Thank you! Yes, this fixed the problem of the original adapter.

If I add more complex transformations using a lambda function:

import os

if "KERAS_BACKEND" not in os.environ:
    # set this to "torch", "tensorflow", or "jax"
    os.environ["KERAS_BACKEND"] = "torch"

import bayesflow as bf
import keras
import numpy as np

def meta(batch_size):
    # batch_size needs to be present but will be ignored here
    # N: number of observation in a dataset
    N = np.random.randint(5, 15)
    return dict(N=N)

def prior():
    # beta: regression coefficients (intercept, slope)
    beta = np.random.normal([2, 0], [3, 1])
    # sigma: residual standard deviation
    sigma = np.random.gamma(1, 1)
    return dict(beta=beta, sigma=sigma)

def likelihood(beta, sigma, N):
    # x: predictor variable
    x = np.random.normal(0, 1, size=N)
    # y: response variable
    y = np.random.normal(beta[0] + beta[1] * x, sigma, size=N)
    return dict(y=y, x=x)

simulator = bf.simulators.make_simulator([prior, likelihood], meta_fn=meta)
adapter = (
    bf.Adapter()
    .broadcast("N", to="x")
    .as_set(["x", "y"])
    .constrain("sigma", lower=0)
    .standardize(exclude=["N"])
    .apply(include="N", forward=lambda n: np.sqrt(n), inverse=lambda n: n**2)
    .concatenate(["beta", "sigma"], into="inference_variables")
    .concatenate(["x", "y"], into="summary_variables")
    .rename("N", "inference_conditions")
)
inference_network = bf.networks.FlowMatching()
summary_network = bf.networks.DeepSet(depth=2)
approximator = bf.ContinuousApproximator(
   inference_network=inference_network,
   summary_network=summary_network,
   adapter=adapter,
)
epochs = 1
num_batches = 1
batch_size = 8
optimizer = keras.optimizers.Adam(learning_rate=5e-4, clipnorm=1.0)
approximator.compile(optimizer=optimizer)
history = approximator.fit(
    epochs=epochs,
    num_batches=num_batches,
    batch_size=batch_size,
    simulator=simulator,
)

approximator.save("approximator.keras")
loaded_approximator = keras.saving.load_model("approximator.keras")

# Set the number of posterior draws you want to get
num_samples = 1000

# Simulate validation data
val_sims = simulator.sample(200)
# exclude parameters from the conditions used as validation data
conditions = {k: v for k, v in val_sims.items() if k != "beta"}

# obtain num_samples samples of the parameter posterior for every valididation dataset
pdraws = loaded_approximator.sample(conditions=conditions, num_samples=num_samples)

I get a different error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[12], line 64
     56 history = approximator.fit(
     57     epochs=epochs,
     58     num_batches=num_batches,
     59     batch_size=batch_size,
     60     simulator=simulator,
     61 )
     63 approximator.save("approximator.keras")
---> 64 loaded_approximator = keras.saving.load_model("approximator.keras")
     66 # Set the number of posterior draws you want to get
     67 num_samples = 1000

File ~/.conda/envs/bayesflow/lib/python3.11/site-packages/keras/src/saving/saving_api.py:189, in load_model(filepath, custom_objects, compile, safe_mode)
    186         is_keras_zip = True
    188 if is_keras_zip or is_keras_dir or is_hf:
--> 189     return saving_lib.load_model(
    190         filepath,
    191         custom_objects=custom_objects,
    192         compile=compile,
    193         safe_mode=safe_mode,
    194     )
    195 if str(filepath).endswith((".h5", ".hdf5")):
    196     return legacy_h5_format.load_model_from_hdf5(
...
    650         )
    651     return python_utils.func_load(inner_config["value"])
    652 if tf is not None and config["class_name"] == "__typespec__":

ValueError: Requested the deserialization of a `lambda` object. This carries a potential risk of arbitrary code execution and thus it is disallowed by default. If you trust the source of the saved model, you can pass `safe_mode=False` to the loading function in order to allow `lambda` loading, or call `keras.config.enable_unsafe_deserialization()`.

If I try with keras.saving.load_model("approximator.keras", safe_mode=False) I get this error:

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[13], line 75
     72 conditions = {k: v for k, v in val_sims.items() if k != "beta"}
     74 # obtain num_samples samples of the parameter posterior for every valididation dataset
---> 75 pdraws = loaded_approximator.sample(conditions=conditions, num_samples=num_samples)

File ~/projects/bayesflow/bayesflow/approximators/continuous_approximator.py:142, in ContinuousApproximator.sample(self, num_samples, conditions, split, **kwargs)
    134 def sample(
    135     self,
    136     *,
   (...)
    140     **kwargs,
    141 ) -> dict[str, np.ndarray]:
--> 142     conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
    143     conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
    144     conditions = {"inference_variables": self._sample(num_samples=num_samples, **conditions, **kwargs)}

File ~/projects/bayesflow/bayesflow/adapters/adapter.py:76, in Adapter.__call__(self, data, inverse, **kwargs)
     73 if inverse:
     74     return self.inverse(data, **kwargs)
---> 76 return self.forward(data, **kwargs)

File ~/projects/bayesflow/bayesflow/adapters/adapter.py:60, in Adapter.forward(self, data, **kwargs)
     57 data = data.copy()
...
     43 )
     44 inference_network = bf.networks.FlowMatching()
     45 summary_network = bf.networks.DeepSet(depth=2)

NameError: name 'np' is not defined

Should I avoid the use of lambda functions and define the function separately?

Thanks again!

5

Yes, please avoid the use of lambda functions, they are not really compatible with the way serialization works. Keras (the backend we use) basically stores names of classes and functions to identify them. So at loading time, it has to match those names to the actual classes and functions again, which does not work for lambdas.

I think it would be nice to have a tutorial notebook explaining the dos and don’ts for serialization, especially with custom functions and classes. Until we get to that, here is some background in case you are interested (based on this Keras guide):

We will take a quick look at the internals to understand which options we have. Let’s start by turning the lambda for the forward transform into a function:

def forward_transform(n):
    return np.sqrt(n)

Keras allows us to view the name it assigns to a given object using keras.saving.get_registered_name, which will be used during serialization. For our function, this gives

>>> keras.saving.get_registered_name(forward_transform)
'forward_transform'

At deserialization, Keras will internally use keras.saving.get_registered_object to turn the name into an object again. Let’s call it with the name we got:

>>> keras.saving.get_registered_object('forward_transform')
None

As we see, Keras does not know an object with that name, so deserialization would fail if we had used forward_transform somewhere. We have two ways to tackle this. First, we can specify which object should be associated with the name 'forward_transform', using the custom_objects parameter which is available in all deserialization functions. Let’s try this:

>>> keras.saving.get_registered_object('forward_transform', custom_objects={'forward_transform': forward_transform})
<function __main__.forward_transform(n)>

This gives the correct function (if we did not change it between storing in loading, which Keras would in general would not be able to detect).

Second, to avoid manually specifying custom objects, we can also register them using the @keras.saving.register_keras_serializable decorator:

@keras.saving.register_keras_serializable(package="custom_package")
def forward_transform(n):
    return np.sqrt(n)

Now we get:

>>> keras.saving.get_registered_name(forward_transform)
'custom_package>forward_transform'
>>> keras.saving.get_registered_object('custom_package>forward_transform')
<function __main__.forward_transform(n)>

Keras returns the correct function without specifying custom_objects. Note that registering is not persistent between sessions, so if you train and store your model, but want to load it at a later time, you have to register the functions again before trying to load the object, as Keras only stores the name and not the function itself.

Usually, we do not have to call get_registered_name and get_registered_object ourselves, as all of this is handled internally. What we have to do is to ensure that we either have registered the functions we use with register_keras_serializable, or manually provide them using custom_objects. The latter is especially useful when we want to load an object and realize that not all functions we used were registered when we stored it.

To go back to the start, for lambda functions we get

>>> keras.saving.get_registered_name(lambda n: np.sqrt(n))
'<lambda>'
>>> keras.saving.get_registered_name(lambda n: n**2)
'<lambda>'

so we just can’t tell them apart when we try to store and load them.

Thanks for bringing this up, we can include a warning if we detect that lambda functions are used in Adapter.apply, and maybe even check whether the functions are registered.

3 Likes