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

**URL:** <https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127>\
**Category:** General\
**Created:** [February 12, 2025, 3:38pm UTC](https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127 "2025-02-12T15:38:16Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![philippreiser](https://avatars.discourse-cdn.com/v4/letter/p/f04885/32.png) [@philippreiser](https://discuss.bayesflow.org/u/philippreiser)\
**Post date:** [February 12, 2025, 3:38pm UTC](https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127/1 "2025-02-12T15:38:16Z")

</div>

Hi!

I am getting started to work with BayesFlow on the `dev` branch (latest commit 8210610). Sorry for all the questions! 😃  
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:

```python
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:

```python
---------------------------------------------------------------------------
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!

---

<div class="post-metadata">

**Author:** ![valentin](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/valentin/32/103_2.png) [@valentin](https://discuss.bayesflow.org/u/valentin)\
**Post date:** [February 13, 2025, 10:27am UTC](https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127/2 "2025-02-13T10:27:53Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![valentin](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/valentin/32/103_2.png) [@valentin](https://discuss.bayesflow.org/u/valentin)\
**Post date:** [February 13, 2025, 1:24pm UTC](https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127/3 "2025-02-13T13:24:25Z")

</div>

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!

---

<div class="post-metadata">

**Author:** ![philippreiser](https://avatars.discourse-cdn.com/v4/letter/p/f04885/32.png) [@philippreiser](https://discuss.bayesflow.org/u/philippreiser)\
**Post date:** [February 13, 2025, 3:55pm UTC](https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127/4 "2025-02-13T15:55:46Z")

</div>

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

If I add more complex transformations using a lambda function:

```python
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:

```python
---------------------------------------------------------------------------
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:

```python
---------------------------------------------------------------------------
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!

---

<div class="post-metadata">

**Author:** ![valentin](https://yyz1.discourse-cdn.com/flex007/user_avatar/discuss.bayesflow.org/valentin/32/103_2.png) [@valentin](https://discuss.bayesflow.org/u/valentin)\
**Post date:** [February 16, 2025, 9:28am UTC](https://discuss.bayesflow.org/t/how-to-correctly-load-a-trained-continuousapproximator-model-to-sample-from-it/127/5 "2025-02-16T09:28:42Z")

</div>

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](https://keras.io/guides/serialization_and_saving/#custom-objects)):

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:

```auto
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

```python
>>> 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:

```auto
>>> 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:

```auto
>>> 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:

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

```

Now we get:

```auto
>>> 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

```auto
>>> 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.