Model portability issue across machines

General
1

Hello,

Together with a collaborator, we are using BayesFlow for amortised posterior estimation and have run into a model portability issue that we have not been able to resolve. We would be very grateful for any guidance.
We have a trained BayesFlow model (BasicWorkflow with CouplingFlow, 12 layers, hidden dim 256, no summary network) that works correctly on Machine A. We are trying to load it on Machine B to run inference on new data without retraining.

Machine A: BayesFlow 2.0.3, Keras 3.10.0, TensorFlow 2.16.1, Apple M1 Max (Metal backend)
Machine B: BayesFlow 2.0.3, Keras 3.10.0, TensorFlow 2.16.2, Ubuntu 26.04 LTS (x86_64, CPU)

When we try to load the complete full_model.keras on Machine B, keras.models.load_model fails because the saved configuration contains standardize=['inference_variables'], which is passed to ContinuousApproximator, but the ContinuousApproximator implementation on Machine B does not accept this argument.

Interestingly, although both machines report BayesFlow 2.0.3, the ContinuousApproximator.__init__ signatures differ: Machine A accepts standardize, whereas Machine B does not.

We therefore tried to bypass the serialization issue by rebuilding the architecture on Machine B and loading only the inference-network weights. However, we found that the resulting architecture has 600,768 trainable parameters, compared with 600,865 on Machine A — a difference of 97 parameters. These 97 parameters appear to correspond to the standardization components (48 means + 48 variances + 1 additional parameter), which are present on Machine A but not correctly reproduced on Machine B despite explicitly setting standardize='all'.

Although the inference-network weights can be loaded, the resulting model produces incorrect posteriors. For example, the posterior for the precision parameter barely moves from the prior on Machine B, whereas on Machine A it is correctly concentrated well above the prior mean.

What we have tried

  • Saving/loading the full model with keras.models.load_model
  • Saving/loading weights only with save_weights / load_weights
  • Saving inference-network weights as NumPy .npz and loading them with set_weights
  • Explicitly setting standardize='all' on Machine B
  • Verifying that all other inputs (feature vectors and normalization statistics) are identical

All approaches either fail during loading or produce incorrect inference results.
Is there a recommended way to save and load a trained BayesFlow model portably across machines with different operating systems and hardware backends? In particular, is there a way to ensure that the standardization layers are serialized and reconstructed consistently, so that we can reproduce on Machine B exactly the trained model from Machine A?

Thank you very much for your help.

Best wishes,
Amanda Lenzi

2

Welcome to Bayesflow discourse! At some point around 2.0.3 we have been moving the (adaptive) standardization feature from adapter to approximator. This is possibly where your issue arises from.

Can you try intalling 2.0.3 from the github (in a new virtual env to not break your existing workflow) and see if it works for that?

3

It works now, thank you very much!