amazing keras fix

Author: LarsKue

bayesflow/approximators/approximator.py

@@ -11,8 +11,8 @@
 
 
 class Approximator(BackendApproximator):
-    def build(self, data_shapes: any) -> None:
-        mock_data = keras.tree.map_structure(keras.ops.zeros, data_shapes)
+    def build(self, data_shapes: dict[str, tuple[int]]) -> None:
+        mock_data = {key: keras.ops.zeros(value) for key, value in data_shapes.items()}
         self.build_from_data(mock_data)
 
     @classmethod
@@ -61,6 +61,9 @@ def build_dataset(
             max_queue_size=max_queue_size,
         )
 
+    def call(self, *args, **kwargs):
+        return self.compute_metrics(*args, **kwargs)
+
     def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = None, **kwargs):
         """
         Trains the approximator on the provided dataset or on-demand data generated from the given simulator.

bayesflow/approximators/continuous_approximator.py

@@ -49,17 +49,23 @@ def __init__(
         self.inference_network = inference_network
         self.summary_network = summary_network
 
-        if standardize == "all":
-            standardize = ["inference_variables", "summary_variables", "inference_conditions"]
-        elif isinstance(standardize, str):
-            standardize = [standardize]
-        elif isinstance(standardize, Sequence):
-            standardize = standardize
-        else:
-            standardize = []
+        # if standardize == "all":
+        #     standardize = ["inference_variables", "summary_variables", "inference_conditions"]
+        # elif isinstance(standardize, str):
+        #     standardize = [standardize]
+        # elif isinstance(standardize, Sequence):
+        #     standardize = standardize
+        # else:
+        #     standardize = []
 
         self.standardize = standardize
-        self.standardize_layers = {s: Standardization() for s in standardize}
+
+        if standardize == "all":
+            # we have to lazily initialize these
+            self.standardize_layers = None
+        else:
+            print("eager init")
+            self.standardize_layers = {s: Standardization(trainable=False) for s in self.standardize}
 
     @classmethod
     def build_adapter(
@@ -121,7 +127,16 @@ def compile(
         return super().compile(*args, **kwargs)
 
     def build_from_data(self, adapted_data: dict[str, any]):
+        if self.standardize == "all":
+            self.standardize = list(adapted_data.keys())
+            self.standardize = ["inference_variables", "summary_variables", "inference_conditions"]
+            self.standardize = list(filter(lambda x: x in adapted_data, self.standardize))
+
+        if self.standardize_layers is None:
+            self.standardize_layers = {s: Standardization(trainable=False) for s in self.standardize}
+
         self.compute_metrics(**filter_kwargs(adapted_data, self.compute_metrics), stage="training")
+
         self.built = True
 
     def compile_from_config(self, config):
@@ -207,11 +222,11 @@ def _compute_summary_metrics(self, summary_variables: Tensor | None, stage: str)
         summary_outputs = summary_metrics.pop("outputs")
         return summary_metrics, summary_outputs
 
-    def _prepare_inference_variables(self, inference_variables, stage):
+    def _prepare_inference_variables(self, inference_variables: Tensor, stage: str) -> Tensor:
         """Helper function to convert inference variables to tensors and optionally standardize them."""
-        inference_variables = keras.tree.map_structure(keras.ops.convert_to_tensor, inference_variables)
         if "inference_variables" in self.standardize:
             inference_variables = self.standardize_layers["inference_variables"](inference_variables, stage=stage)
+
         return inference_variables
 
     def _combine_conditions(

bayesflow/networks/standardization/standardization.py

@@ -4,12 +4,12 @@
 
 from bayesflow.types import Tensor, Shape
 from bayesflow.utils.serialization import serialize, deserialize, serializable
-from bayesflow.utils import expand_left_as
+from bayesflow.utils import expand_left_as, layer_kwargs
 
 
 @serializable("bayesflow.networks")
 class Standardization(keras.Layer):
-    def __init__(self, momentum: float = 0.95, epsilon: float = 1e-6):
+    def __init__(self, momentum: float = 0.95, epsilon: float = 1e-6, **kwargs):
         """
         Initializes a Standardization layer that will keep track of the running mean and
         running standard deviation across a batch of tensors.
@@ -23,14 +23,17 @@ def __init__(self, momentum: float = 0.95, epsilon: float = 1e-6):
         epsilon: float, optional
             Stability parameter to avoid division by zero.
         """
-        super().__init__()
+        super().__init__(**layer_kwargs(kwargs))
 
         self.momentum = momentum
         self.epsilon = epsilon
         self.moving_mean = None
         self.moving_std = None
 
     def build(self, input_shape: Shape):
+        if self.built:
+            return
+
         self.moving_mean = self.add_weight(shape=(input_shape[-1],), initializer="zeros", trainable=False)
         self.moving_std = self.add_weight(shape=(input_shape[-1],), initializer="ones", trainable=False)
 

tests/test_approximators/test_serialize_deserialize.py

@@ -0,0 +1,26 @@
+import keras
+import numpy as np
+
+from ..utils import assert_models_equal
+
+
+def test_serialize_deserialize_continuous_approximator(tmp_path, continuous_approximator):
+    sample_data = {
+        "mean": np.zeros((32, 10, 2)),
+        "std": np.ones((32, 10, 1)),
+        "x": np.random.standard_normal((32, 10, 2)),
+    }
+
+    sample_data = continuous_approximator.adapter(sample_data)
+
+    continuous_approximator.build_from_data(sample_data)
+
+    keras.saving.save_model(continuous_approximator, tmp_path / "model.keras")
+    loaded = keras.saving.load_model(tmp_path / "model.keras")
+    assert_models_equal(continuous_approximator, loaded)
+
+    # serialized = serialize(continuous_approximator)
+    # deserialized = deserialize(serialized)
+    # reserialized = serialize(deserialized)
+    #
+    # assert serialized == reserialized

tests/test_two_moons/test_two_moons.py

@@ -23,7 +23,8 @@ def test_fit(approximator, train_dataset, validation_dataset, batch_size):
 
     mock_data = train_dataset[0]
     mock_data = keras.tree.map_structure(keras.ops.convert_to_tensor, mock_data)
-    approximator.build_from_data(mock_data)
+    mock_data_shapes = keras.tree.map_structure(keras.ops.shape, mock_data)
+    approximator.build(mock_data_shapes)
 
     untrained_weights = copy.deepcopy(approximator.weights)
     untrained_metrics = approximator.evaluate(validation_dataset, return_dict=True)