🚀 Merge pull request #413 from bayesflow-org/allow-networks

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

Author: LarsKue

bayesflow/adapters/adapter.py

@@ -2,11 +2,7 @@
 
 import numpy as np
 
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+from bayesflow.utils.serialization import deserialize, serialize, serializable
 
 from .transforms import (
     AsSet,
@@ -33,7 +29,7 @@
 from .transforms.filter_transform import Predicate
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Adapter(MutableSequence[Transform]):
     """
     Defines an adapter to apply various transforms to data.
@@ -74,10 +70,14 @@ def create_default(inference_variables: Sequence[str]) -> "Adapter":
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Adapter":
-        return cls(transforms=deserialize(config["transforms"], custom_objects))
+        return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_config(self) -> dict:
-        return {"transforms": serialize(self.transforms)}
+        config = {
+            "transforms": self.transforms,
+        }
+
+        return serialize(config)
 
     def forward(self, data: dict[str, any], **kwargs) -> dict[str, np.ndarray]:
         """Apply the transforms in the forward direction.

bayesflow/adapters/transforms/as_set.py

@@ -1,10 +1,11 @@
-from keras.saving import register_keras_serializable as serializable
 import numpy as np
 
+from bayesflow.utils.serialization import serializable
+
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class AsSet(ElementwiseTransform):
     """The `.as_set(["x", "y"])` transform indicates that both `x` and `y` are treated as sets.
 
@@ -33,9 +34,5 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
         return data
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "AsSet":
-        return cls()
-
     def get_config(self) -> dict:
         return {}

bayesflow/adapters/transforms/as_time_series.py

@@ -1,10 +1,11 @@
 import numpy as np
-from keras.saving import register_keras_serializable as serializable
+
+from bayesflow.utils.serialization import serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class AsTimeSeries(ElementwiseTransform):
     """The `.as_time_series` transform can be used to indicate that variables shall be treated as time series.
 
@@ -29,9 +30,5 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
 
         return data
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "AsTimeSeries":
-        return cls()
-
     def get_config(self) -> dict:
         return {}

bayesflow/adapters/transforms/broadcast.py

@@ -1,16 +1,12 @@
 from collections.abc import Sequence
 import numpy as np
 
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+from bayesflow.utils.serialization import serialize, serializable
 
 from .transform import Transform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Broadcast(Transform):
     """
     Broadcasts arrays or scalars to the shape of a given other array.
@@ -96,31 +92,15 @@ def __init__(
         self.exclude = exclude
         self.squeeze = squeeze
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "Broadcast":
-        # Deserialize turns tuples to lists, undo it if necessary
-        exclude = deserialize(config["exclude"], custom_objects)
-        exclude = tuple(exclude) if isinstance(exclude, list) else exclude
-        expand = deserialize(config["expand"], custom_objects)
-        expand = tuple(expand) if isinstance(expand, list) else expand
-        squeeze = deserialize(config["squeeze"], custom_objects)
-        squeeze = tuple(squeeze) if isinstance(squeeze, list) else squeeze
-        return cls(
-            keys=deserialize(config["keys"], custom_objects),
-            to=deserialize(config["to"], custom_objects),
-            expand=expand,
-            exclude=exclude,
-            squeeze=squeeze,
-        )
-
     def get_config(self) -> dict:
-        return {
-            "keys": serialize(self.keys),
-            "to": serialize(self.to),
-            "expand": serialize(self.expand),
-            "exclude": serialize(self.exclude),
-            "squeeze": serialize(self.squeeze),
+        config = {
+            "keys": self.keys,
+            "to": self.to,
+            "expand": self.expand,
+            "exclude": self.exclude,
+            "squeeze": self.squeeze,
         }
+        return serialize(config)
 
     # noinspection PyMethodOverriding
     def forward(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, np.ndarray]:

bayesflow/adapters/transforms/concatenate.py

@@ -1,16 +1,13 @@
 from collections.abc import Sequence
 
 import numpy as np
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+
+from bayesflow.utils.serialization import serialize, serializable
 
 from .transform import Transform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Concatenate(Transform):
     """Concatenate multiple arrays into a new key. Used to specify how data variables should be treated by the network.
 
@@ -35,29 +32,21 @@ class Concatenate(Transform):
     )
     """
 
-    def __init__(self, keys: Sequence[str], *, into: str, axis: int = -1, _indices: list | None = None):
+    def __init__(self, keys: Sequence[str], *, into: str, axis: int = -1, indices: list | None = None):
         self.keys = keys
         self.into = into
         self.axis = axis
 
-        self.indices = _indices
-
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "Concatenate":
-        return cls(
-            keys=deserialize(config["keys"], custom_objects),
-            into=deserialize(config["into"], custom_objects),
-            axis=deserialize(config["axis"], custom_objects),
-            _indices=deserialize(config["indices"], custom_objects),
-        )
+        self.indices = indices
 
     def get_config(self) -> dict:
-        return {
-            "keys": serialize(self.keys),
-            "into": serialize(self.into),
-            "axis": serialize(self.axis),
-            "indices": serialize(self.indices),
+        config = {
+            "keys": self.keys,
+            "into": self.into,
+            "axis": self.axis,
+            "indices": self.indices,
         }
+        return serialize(config)
 
     def forward(self, data: dict[str, any], *, strict: bool = True, **kwargs) -> dict[str, any]:
         if not strict and self.indices is None:

bayesflow/adapters/transforms/constrain.py

@@ -1,8 +1,6 @@
-from keras.saving import (
-    register_keras_serializable as serializable,
-)
 import numpy as np
 
+from bayesflow.utils.serialization import serializable, serialize
 from bayesflow.utils.numpy_utils import (
     inverse_sigmoid,
     inverse_softplus,
@@ -13,7 +11,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Constrain(ElementwiseTransform):
     """
     Constrains neural network predictions of a data variable to specified bounds.
@@ -163,18 +161,15 @@ def unconstrain(x):
             case other:
                 raise ValueError(f"Unsupported value for 'inclusive': {other!r}.")
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "Constrain":
-        return cls(**config)
-
     def get_config(self) -> dict:
-        return {
+        config = {
             "lower": self.lower,
             "upper": self.upper,
             "method": self.method,
             "inclusive": self.inclusive,
             "epsilon": self.epsilon,
         }
+        return serialize(config)
 
     def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         # forward means data space -> network space, so unconstrain the data

bayesflow/adapters/transforms/convert_dtype.py

@@ -1,14 +1,11 @@
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
 import numpy as np
 
+from bayesflow.utils.serialization import serializable, serialize
+
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class ConvertDType(ElementwiseTransform):
     """
     Default transform used to convert all floats from float64 to float32 to be in line with keras framework.
@@ -27,21 +24,15 @@ def __init__(self, from_dtype: str, to_dtype: str):
         self.from_dtype = from_dtype
         self.to_dtype = to_dtype
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "ConvertDType":
-        return cls(
-            from_dtype=deserialize(config["from_dtype"], custom_objects),
-            to_dtype=deserialize(config["to_dtype"], custom_objects),
-        )
-
     def get_config(self) -> dict:
-        return {
-            "from_dtype": serialize(self.from_dtype),
-            "to_dtype": serialize(self.to_dtype),
+        config = {
+            "from_dtype": self.from_dtype,
+            "to_dtype": self.to_dtype,
         }
+        return serialize(config)
 
     def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        return data.astype(self.to_dtype)
+        return data.astype(self.to_dtype, copy=False)
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        return data.astype(self.from_dtype)
+        return data.astype(self.from_dtype, copy=False)

bayesflow/adapters/transforms/drop.py

@@ -1,15 +1,11 @@
 from collections.abc import Sequence
 
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+from bayesflow.utils.serialization import serializable, serialize
 
 from .transform import Transform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class Drop(Transform):
     """
     Transform to drop variables from further calculation.
@@ -37,12 +33,8 @@ class Drop(Transform):
     def __init__(self, keys: Sequence[str]):
         self.keys = keys
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "Drop":
-        return cls(keys=deserialize(config["keys"], custom_objects))
-
     def get_config(self) -> dict:
-        return {"keys": serialize(self.keys)}
+        return serialize({"keys": self.keys})
 
     def forward(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         # no strict version because there is no requirement for the keys to be present

bayesflow/adapters/transforms/elementwise_transform.py

@@ -1,8 +1,9 @@
-from keras.saving import register_keras_serializable as serializable
 import numpy as np
 
+from bayesflow.utils.serialization import serializable, deserialize
 
-@serializable(package="bayesflow.adapters")
+
+@serializable
 class ElementwiseTransform:
     """Base class on which other transforms are based"""
 
@@ -13,8 +14,8 @@ def __call__(self, data: np.ndarray, inverse: bool = False, **kwargs) -> np.ndar
         return self.forward(data, **kwargs)
 
     @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "ElementwiseTransform":
-        raise NotImplementedError
+    def from_config(cls, config: dict, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_config(self) -> dict:
         raise NotImplementedError

bayesflow/adapters/transforms/expand_dims.py

@@ -1,14 +1,11 @@
 import numpy as np
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
-)
+
+from bayesflow.utils.serialization import serializable, serialize
 
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable(package="bayesflow.adapters")
+@serializable
 class ExpandDims(ElementwiseTransform):
     """
     Expand the shape of an array.
@@ -51,16 +48,8 @@ def __init__(self, *, axis: int | tuple):
         super().__init__()
         self.axis = axis
 
-    @classmethod
-    def from_config(cls, config: dict, custom_objects=None) -> "ExpandDims":
-        return cls(
-            axis=deserialize(config["axis"], custom_objects),
-        )
-
     def get_config(self) -> dict:
-        return {
-            "axis": serialize(self.axis),
-        }
+        return serialize({"axis": self.axis})
 
     def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         return np.expand_dims(data, axis=self.axis)