Improve serialization and allow networks to be passed directly to most models

bayesflow/distributions/__init__.py

@@ -8,6 +8,8 @@
 from .diagonal_normal import DiagonalNormal
 from .diagonal_student_t import DiagonalStudentT
 
+from .find_distribution import find_distribution
+
 from ..utils._docs import _add_imports_to_all
 
 _add_imports_to_all(include_modules=[])

bayesflow/distributions/distribution.py

@@ -1,12 +1,12 @@
 import keras
 
 from bayesflow.types import Shape, Tensor
-from bayesflow.utils import keras_kwargs
+from bayesflow.utils import layer_kwargs
 
 
 class Distribution(keras.Layer):
     def __init__(self, **kwargs):
-        super().__init__(**keras_kwargs(kwargs))
+        super().__init__(**layer_kwargs(kwargs))
 
     def call(self, samples: Tensor) -> Tensor:
         return keras.ops.exp(self.log_prob(samples))
@@ -16,3 +16,6 @@
 
     def sample(self, batch_shape: Shape) -> Tensor:
         raise NotImplementedError
+
+    def compute_output_shape(self, input_shape: Shape) -> Shape:
+        return keras.ops.shape(self.sample(input_shape[0:1]))

bayesflow/utils/dispatch/find_distribution.py → bayesflow/distributions/find_distribution.py

@@ -1,5 +1,7 @@
 from functools import singledispatch
 
+from bayesflow.distributions import Distribution
+
 
 @singledispatch
 def find_distribution(arg, **kwargs):
@@ -24,3 +26,8 @@ def _(name: str, *args, **kwargs):
 @find_distribution.register
 def _(none: None, *args, **kwargs):
     return None
+
+
+@find_distribution.register
+def _(distribution: Distribution, *args, **kwargs):
+    return distribution

bayesflow/experimental/cif/conditional_gaussian.py

@@ -4,7 +4,7 @@
 from bayesflow.networks.mlp import MLP
 
 from bayesflow.types import Shape, Tensor
-from bayesflow.utils import keras_kwargs
+from bayesflow.utils import layer_kwargs
 
 
 @register_keras_serializable(package="bayesflow.networks.cif")
@@ -32,7 +32,7 @@
             The MLP activation function
         """
 
-        super().__init__(**keras_kwargs(kwargs))
+        super().__init__(**layer_kwargs(kwargs))
         self.means = MLP([width] * depth, activation=activation)
         self.stds = MLP([width] * depth, activation=activation)
         self.output_projector = keras.layers.Dense(None)

bayesflow/experimental/continuous_time_consistency_model.py

@@ -1,29 +1,28 @@
 import keras
 from keras import ops
-from keras.saving import (
-    register_keras_serializable,
-)
 
 import numpy as np
 
+import warnings
+
+from bayesflow.networks import MLP
 from bayesflow.types import Tensor
 from bayesflow.utils import (
     jvp,
     concatenate_valid,
     find_network,
-    keras_kwargs,
     expand_right_as,
     expand_right_to,
-    serialize_value_or_type,
-    deserialize_value_or_type,
+    model_kwargs,
 )
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 
 from bayesflow.networks import InferenceNetwork
 from bayesflow.networks.embeddings import FourierEmbedding
 
 
-@register_keras_serializable(package="bayesflow.networks")
+@serializable
 class ContinuousTimeConsistencyModel(InferenceNetwork):
     """Implements an sCM (simple, stable, and scalable Consistency Model)
     with continous-time Consistency Training (CT) as described in [1].
@@ -40,8 +39,10 @@
 
     def __init__(
         self,
-        subnet: str | type = "mlp",
+        subnet: str | keras.Layer = "mlp",
         sigma_data: float = 1.0,
+        subnet_kwargs: dict[str, any] = None,
+        embedding_kwargs: dict[str, any] = None,
         **kwargs,
     ):
         """Creates an instance of an sCM to be used for consistency training (CT).
@@ -53,39 +54,52 @@
             instantiated using subnet_kwargs.
         sigma_data    : float, optional, default: 1.0
             Standard deviation of the target distribution
-        **kwargs      : dict, optional, default: {}
-            Additional keyword arguments, such as
+        **kwargs
+            Additional keyword arguments to the layer.
         """
-        super().__init__(base_distribution="normal", **keras_kwargs(kwargs))
+        super().__init__(base_distribution="normal", **kwargs)
+
+        if subnet_kwargs:
+            warnings.warn(
+                "Using `subnet_kwargs` is deprecated."
+                "Instead, instantiate the network yourself and pass the arguments directly.",
+                DeprecationWarning,
+            )
+
+        subnet_kwargs = subnet_kwargs or {}
 
-        self.subnet = find_network(subnet, **kwargs.get("subnet_kwargs", {}))
-        self.subnet_projector = keras.layers.Dense(units=None, bias_initializer="zeros", kernel_initializer="zeros")
+        self.subnet = find_network(subnet, **subnet_kwargs)
+        self.subnet_projector = keras.layers.Dense(
+            units=None, bias_initializer="zeros", kernel_initializer="zeros", name="subnet_projector"
+        )
 
-        self.weight_fn = find_network("mlp", widths=(256,), dropout=0.0)
-        self.weight_fn_projector = keras.layers.Dense(units=1, bias_initializer="zeros", kernel_initializer="zeros")
+        self.weight_fn = MLP([256], dropout=0.0)
+        self.weight_fn_projector = keras.layers.Dense(
+            units=1, bias_initializer="zeros", kernel_initializer="zeros", name="weight_fn_projector"
+        )
 
-        self.time_emb = FourierEmbedding(**kwargs.get("embedding_kwargs", {}))
+        embedding_kwargs = embedding_kwargs or {}
+        self.time_emb = FourierEmbedding(**embedding_kwargs)
         self.time_emb_dim = self.time_emb.embed_dim
 
         self.sigma_data = sigma_data
 
         self.seed_generator = keras.random.SeedGenerator()
 
-        # serialization: store all parameters necessary to call __init__
-        self.config = {
-            "sigma_data": sigma_data,
-            **kwargs,
-        }
-        self.config = serialize_value_or_type(self.config, "subnet", subnet)
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_config(self):
         base_config = super().get_config()
-        return base_config | self.config
+        base_config = model_kwargs(base_config)
 
-    @classmethod
-    def from_config(cls, config):
-        config = deserialize_value_or_type(config, "subnet")
-        return cls(**config)
+        config = {
+            "subnet": self.subnet,
+            "sigma_data": self.sigma_data,
+        }
+
+        return base_config | serialize(config)
 
     def _discretize_time(self, num_steps: int, rho: float = 3.5, **kwargs):
         t = np.linspace(0.0, np.pi / 2, num_steps)
@@ -206,7 +220,9 @@
         out = ops.cos(t) * x - ops.sin(t) * self.sigma_data * f
         return out
 
-    def compute_metrics(self, x: Tensor, conditions: Tensor = None, stage: str = "training") -> dict[str, Tensor]:
+    def compute_metrics(
+        self, x: Tensor, conditions: Tensor = None, stage: str = "training", **kwargs
+    ) -> dict[str, Tensor]:
         base_metrics = super().compute_metrics(x, conditions=conditions, stage=stage)
 
         # $# Implements Algorithm 1 from [1]

bayesflow/experimental/free_form_flow/free_form_flow.py

@@ -1,24 +1,25 @@
 import keras
 from keras import ops
-from keras.saving import register_keras_serializable as serializable
 
+import warnings
+
+from bayesflow.distributions import Distribution
 from bayesflow.types import Tensor
 from bayesflow.utils import (
-    find_network,
-    keras_kwargs,
     concatenate_valid,
+    find_network,
     jacobian,
     jvp,
+    model_kwargs,
     vjp,
-    serialize_value_or_type,
-    deserialize_value_or_type,
     weighted_mean,
 )
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from bayesflow.networks import InferenceNetwork
 
 
-@serializable(package="networks.free_form_flow")
+@serializable
 class FreeFormFlow(InferenceNetwork):
     """Implements a dimensionality-preserving Free-form Flow.
     Incorporates ideas from [1-2].
@@ -53,10 +54,12 @@
     def __init__(
         self,
         beta: float = 50.0,
-        encoder_subnet: str | type = "mlp",
-        decoder_subnet: str | type = "mlp",
-        base_distribution: str = "normal",
+        encoder_subnet: str | keras.Layer = "mlp",
+        decoder_subnet: str | keras.Layer = "mlp",
+        base_distribution: str | Distribution = "normal",
         hutchinson_sampling: str = "qr",
+        encoder_subnet_kwargs: dict = None,
+        decoder_subnet_kwargs: dict = None,
         **kwargs,
     ):
         """Creates an instance of a Free-form Flow.
@@ -80,54 +83,48 @@
         **kwargs              : dict, optional, default: {}
             Additional keyword arguments
         """
-        super().__init__(base_distribution=base_distribution, **keras_kwargs(kwargs))
+        super().__init__(base_distribution, **kwargs)
 
-        if encoder_subnet == "mlp":
-            encoder_subnet_kwargs = FreeFormFlow.ENCODER_MLP_DEFAULT_CONFIG.copy()
-            encoder_subnet_kwargs.update(kwargs.get("encoder_subnet_kwargs", {}))
-        else:
-            encoder_subnet_kwargs = kwargs.get("encoder_subnet_kwargs", {})
+        if encoder_subnet_kwargs or decoder_subnet_kwargs:
+            warnings.warn(
+                "Using `subnet_kwargs` is deprecated."
+                "Instead, instantiate the network yourself and pass the arguments directly.",
+                DeprecationWarning,
+            )
 
-        self.encoder_subnet = find_network(encoder_subnet, **encoder_subnet_kwargs)
-        self.encoder_projector = keras.layers.Dense(units=None, bias_initializer="zeros", kernel_initializer="zeros")
+        encoder_subnet_kwargs = encoder_subnet_kwargs or {}
+        decoder_subnet_kwargs = decoder_subnet_kwargs or {}
+
+        if encoder_subnet == "mlp":
+            encoder_subnet_kwargs = FreeFormFlow.ENCODER_MLP_DEFAULT_CONFIG.copy() | encoder_subnet_kwargs
 
         if decoder_subnet == "mlp":
-            decoder_subnet_kwargs = FreeFormFlow.DECODER_MLP_DEFAULT_CONFIG.copy()
-            decoder_subnet_kwargs.update(kwargs.get("decoder_subnet_kwargs", {}))
-        else:
-            decoder_subnet_kwargs = kwargs.get("decoder_subnet_kwargs", {})
+            decoder_subnet_kwargs = FreeFormFlow.DECODER_MLP_DEFAULT_CONFIG.copy() | decoder_subnet_kwargs
+
+        self.encoder_subnet = find_network(encoder_subnet, **encoder_subnet_kwargs)
+        self.encoder_projector = keras.layers.Dense(
+            units=None, bias_initializer="zeros", kernel_initializer="zeros", name="encoder_projector"
+        )
 
         self.decoder_subnet = find_network(decoder_subnet, **decoder_subnet_kwargs)
-        self.decoder_projector = keras.layers.Dense(units=None, bias_initializer="zeros", kernel_initializer="zeros")
+        self.decoder_projector = keras.layers.Dense(
+            units=None, bias_initializer="zeros", kernel_initializer="zeros", name="decoder_projector"
+        )
 
         self.hutchinson_sampling = hutchinson_sampling
         self.beta = beta
 
         self.seed_generator = keras.random.SeedGenerator()
 
-        # serialization: store all parameters necessary to call __init__
-        self.config = {
-            "beta": beta,
-            "base_distribution": base_distribution,
-            "hutchinson_sampling": hutchinson_sampling,
-            **kwargs,
-        }
-        self.config = serialize_value_or_type(self.config, "encoder_subnet", encoder_subnet)
-        self.config = serialize_value_or_type(self.config, "decoder_subnet", decoder_subnet)
-
-    def get_config(self):
-        base_config = super().get_config()
-        return base_config | self.config
-
-    @classmethod
-    def from_config(cls, config):
-        config = deserialize_value_or_type(config, "encoder_subnet")
-        config = deserialize_value_or_type(config, "decoder_subnet")
-        return cls(**config)
-
     # noinspection PyMethodOverriding
     def build(self, xz_shape, conditions_shape=None):
-        super().build(xz_shape)
+        if self.built:
+            # building when the network is already built can cause issues with serialization
+            # see https://github.com/keras-team/keras/issues/21147
+            return
+
+        self.base_distribution.build(xz_shape)
+
         self.encoder_projector.units = xz_shape[-1]
         self.decoder_projector.units = xz_shape[-1]
 
@@ -142,11 +139,11 @@
         self.encoder_subnet.build(input_shape)
         self.decoder_subnet.build(input_shape)
 
-        input_shape = self.encoder_subnet.compute_output_shape(input_shape)
-        self.encoder_projector.build(input_shape)
+        output_shape = self.encoder_subnet.compute_output_shape(input_shape)
+        self.encoder_projector.build(output_shape)
 
-        input_shape = self.decoder_subnet.compute_output_shape(input_shape)
-        self.decoder_projector.build(input_shape)
+        output_shape = self.decoder_subnet.compute_output_shape(input_shape)
+        self.decoder_projector.build(output_shape)
 
     def _forward(
         self, x: Tensor, conditions: Tensor = None, density: bool = False, training: bool = False, **kwargs
@@ -244,3 +241,22 @@
         loss = weighted_mean(losses, sample_weight)
 
         return base_metrics | {"loss": loss}
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))
+
+    def get_config(self):
+        base_config = super().get_config()
+        base_config = model_kwargs(base_config)
+
+        config = {
+            "beta": self.beta,
+            "encoder_subnet": self.encoder_subnet,
+            "decoder_subnet": self.decoder_subnet,
+            "base_distribution": self.base_distribution,
+            "hutchinson_sampling": self.hutchinson_sampling,
+            # we do not need to store subnet_kwargs
+        }
+
+        return base_config | serialize(config)

bayesflow/experimental/improved_mlp/__init__.py

@@ -0,0 +1 @@
+from .mlp import MLP