Major refactor: model to layer

Author: stefanradev93

bayesflow/networks/consistency_models/consistency_model.py

@@ -6,7 +6,7 @@
 import warnings
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, model_kwargs, weighted_mean
+from bayesflow.utils import find_network, layer_kwargs, weighted_mean
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from ..inference_network import InferenceNetwork
@@ -118,7 +118,7 @@ def from_config(cls, config, custom_objects=None):
 
     def get_config(self):
         base_config = super().get_config()
-        base_config = model_kwargs(base_config)
+        base_config = layer_kwargs(base_config)
 
         config = {
             "total_steps": self.total_steps,
@@ -128,6 +128,7 @@ def get_config(self):
             "eps": self.eps,
             "s0": self.s0,
             "s1": self.s1,
+            # we do not need to store subnet_kwargs
         }
 
         return base_config | serialize(config)

bayesflow/networks/coupling_flow/actnorm.py

@@ -1,13 +1,12 @@
 from keras import ops
 
-from keras.saving import register_keras_serializable as serializable
-
 from bayesflow.types import Shape, Tensor
+from bayesflow.utils.serialization import serializable
 
 from .invertible_layer import InvertibleLayer
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class ActNorm(InvertibleLayer):
     """Implements an Activation Normalization (ActNorm) Layer. Activation Normalization is learned invertible
     normalization, using a scale (s) and a bias (b) vector::

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -3,7 +3,7 @@
 from bayesflow.types import Tensor
 from bayesflow.utils import (
     find_permutation,
-    model_kwargs,
+    layer_kwargs,
     weighted_mean,
 )
 from bayesflow.utils.serialization import deserialize, serializable, serialize
@@ -131,7 +131,7 @@ def from_config(cls, config, custom_objects=None):
 
     def get_config(self):
         base_config = super().get_config()
-        base_config = model_kwargs(base_config)
+        base_config = layer_kwargs(base_config)
 
         config = {
             "subnet": self.subnet,

bayesflow/networks/coupling_flow/couplings/dual_coupling.py

@@ -3,7 +3,9 @@
 from bayesflow.utils import model_kwargs
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 from bayesflow.types import Tensor
+
 from .single_coupling import SingleCoupling
+
 from ..invertible_layer import InvertibleLayer
 
 

bayesflow/networks/deep_set/deep_set.py

@@ -3,11 +3,11 @@
 import keras
 
 from bayesflow.types import Tensor
-from bayesflow.utils import model_kwargs
-from bayesflow.utils.serialization import deserialize, serializable, serialize
+from bayesflow.utils.serialization import serializable
+
+from .equivariant_layer import EquivariantLayer
+from .invariant_layer import InvariantLayer
 
-from .equivariant_module import EquivariantModule
-from .invariant_module import InvariantModule
 from ..summary_network import SummaryNetwork
 
 
@@ -88,7 +88,7 @@ def __init__(
         # Stack of equivariant modules for a many-to-many learnable transformation
         self.equivariant_modules = []
         for _ in range(depth):
-            equivariant_module = EquivariantModule(
+            equivariant_module = EquivariantLayer(
                 mlp_widths_equivariant=mlp_widths_equivariant,
                 mlp_widths_invariant_inner=mlp_widths_invariant_inner,
                 mlp_widths_invariant_outer=mlp_widths_invariant_outer,
@@ -97,37 +97,24 @@ def __init__(
                 spectral_normalization=spectral_normalization,
                 dropout=dropout,
                 pooling=inner_pooling,
-                name="equivariant_module",
             )
             self.equivariant_modules.append(equivariant_module)
 
         # Invariant module for a many-to-one transformation
-        self.invariant_module = InvariantModule(
+        self.invariant_module = InvariantLayer(
             mlp_widths_inner=mlp_widths_invariant_last,
             mlp_widths_outer=mlp_widths_invariant_last,
             activation=activation,
             kernel_initializer=kernel_initializer,
             dropout=dropout,
             pooling=output_pooling,
             spectral_normalization=spectral_normalization,
-            name="invariant_module",
         )
 
         # Output linear layer to project set representation down to "summary_dim" learned summary statistics
         self.output_projector = keras.layers.Dense(summary_dim, activation="linear", name="output_projector")
 
         self.summary_dim = summary_dim
-        self.depth = depth
-        self.inner_pooling = inner_pooling
-        self.output_pooling = output_pooling
-        self.mlp_widths_equivariant = mlp_widths_equivariant
-        self.mlp_widths_invariant_inner = mlp_widths_invariant_inner
-        self.mlp_widths_invariant_outer = mlp_widths_invariant_outer
-        self.mlp_widths_invariant_last = mlp_widths_invariant_last
-        self.activation = activation
-        self.kernel_initializer = kernel_initializer
-        self.dropout = dropout
-        self.spectral_normalization = spectral_normalization
 
     def call(self, x: Tensor, training: bool = False, **kwargs) -> Tensor:
         """
@@ -161,28 +148,3 @@ def call(self, x: Tensor, training: bool = False, **kwargs) -> Tensor:
         x = self.invariant_module(x, training=training)
 
         return self.output_projector(x)
-
-    @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 = {
-            "summary_dim": self.summary_dim,
-            "depth": self.depth,
-            "inner_pooling": self.inner_pooling,
-            "output_pooling": self.output_pooling,
-            "mlp_widths_equivariant": self.mlp_widths_equivariant,
-            "mlp_widths_invariant_inner": self.mlp_widths_invariant_inner,
-            "mlp_widths_invariant_outer": self.mlp_widths_invariant_outer,
-            "mlp_widths_invariant_last": self.mlp_widths_invariant_last,
-            "activation": self.activation,
-            "kernel_initializer": self.kernel_initializer,
-            "dropout": self.dropout,
-            "spectral_normalization": self.spectral_normalization,
-        }
-
-        return base_config | serialize(config)

bayesflow/networks/deep_set/equivariant_module.py renamed to bayesflow/networks/deep_set/equivariant_layer.py

@@ -4,17 +4,17 @@
 from keras import ops, layers
 
 from bayesflow.types import Tensor
-from bayesflow.utils import model_kwargs
+from bayesflow.utils import layer_kwargs
 from bayesflow.utils.decorators import sanitize_input_shape
 from bayesflow.utils.serialization import serializable
 
 from ..mlp import MLP
 
-from .invariant_module import InvariantModule
+from .invariant_layer import InvariantLayer
 
 
 @serializable
-class EquivariantModule(keras.Model):
+class EquivariantLayer(keras.Layer):
     """Implements an equivariant module performing an equivariant transform.
 
     For details and justification, see:
@@ -72,32 +72,30 @@ def __init__(
             Whether to apply spectral normalization to stabilize training. Default is False.
         """
 
-        super().__init__(**model_kwargs(kwargs))
+        super().__init__(**layer_kwargs(kwargs))
 
         # Invariant module to increase expressiveness by concatenating outputs to each set member
-        self.invariant_module = InvariantModule(
+        self.invariant_module = InvariantLayer(
             mlp_widths_inner=mlp_widths_invariant_inner,
             mlp_widths_outer=mlp_widths_invariant_outer,
             activation=activation,
             kernel_initializer=kernel_initializer,
             dropout=dropout,
             pooling=pooling,
             spectral_normalization=spectral_normalization,
-            name="invariant_module",
         )
 
         # Fully connected net + residual connection for an equivariant transform applied to each set member
-        self.input_projector = layers.Dense(mlp_widths_equivariant[-1], name="input_projector")
+        self.input_projector = layers.Dense(mlp_widths_equivariant[-1])
         self.equivariant_fc = MLP(
             mlp_widths_equivariant,
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,
             spectral_normalization=spectral_normalization,
-            name="equivariant_fc",
         )
 
-        self.layer_norm = layers.LayerNormalization(name="layer_norm") if layer_norm else None
+        self.layer_norm = layers.LayerNormalization() if layer_norm else None
 
     @sanitize_input_shape
     def build(self, input_shape):

bayesflow/networks/deep_set/invariant_module.py renamed to bayesflow/networks/deep_set/invariant_layer.py

@@ -3,7 +3,7 @@
 import keras
 
 from bayesflow.types import Tensor
-from bayesflow.utils import model_kwargs
+from bayesflow.utils import layer_kwargs
 from bayesflow.utils import find_pooling
 from bayesflow.utils.decorators import sanitize_input_shape
 from bayesflow.utils.serialization import serializable
@@ -12,7 +12,7 @@
 
 
 @serializable
-class InvariantModule(keras.Model):
+class InvariantLayer(keras.Layer):
     """Implements an invariant module performing a permutation-invariant transform.
 
     For details and rationale, see:
@@ -64,7 +64,7 @@ def __init__(
             Whether to apply spectral normalization to stabilize training. Default is False.
         """
 
-        super().__init__(**model_kwargs(kwargs))
+        super().__init__(**layer_kwargs(kwargs))
 
         # Inner fully connected net for sum decomposition: inner( pooling( inner(set) ) )
         self.inner_fc = MLP(
@@ -88,11 +88,6 @@ def __init__(
             pooling_kwargs = {}
 
         self.pooling_layer = find_pooling(pooling, **pooling_kwargs)
-        self.pooling_layer.name = f"{pooling}_pooling"
-
-    @sanitize_input_shape
-    def build(self, input_shape):
-        self.call(keras.ops.zeros(input_shape))
 
     def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         """Performs the forward pass of a learnable invariant transform.
@@ -114,3 +109,7 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         set_summary = self.pooling_layer(set_summary, training=training)
         set_summary = self.outer_fc(set_summary, training=training)
         return set_summary
+
+    @sanitize_input_shape
+    def build(self, input_shape):
+        self.call(keras.ops.zeros(input_shape))

bayesflow/networks/embeddings/fourier_embedding.py

@@ -2,12 +2,12 @@
 
 import keras
 from keras import ops
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
+from bayesflow.utils.serialization import serializable
 
 
-@serializable(package="bayesflow.networks")
+@serializable
 class FourierEmbedding(keras.Layer):
     """Implements a Fourier projection with normally distributed frequencies."""
 

bayesflow/networks/embeddings/recurrent_embedding.py

@@ -1,11 +1,12 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
 from bayesflow.utils import expand_tile
 
+from bayesflow.utils.serialization import serializable
 
-@serializable(package="bayesflow.networks")
+
+@serializable
 class RecurrentEmbedding(keras.Layer):
     """Implements a recurrent network for flexibly embedding time vectors."""
 

bayesflow/networks/embeddings/time2vec.py

@@ -1,11 +1,11 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
 from bayesflow.utils import expand_tile
+from bayesflow.utils.serialization import serializable
 
 
-@serializable(package="bayesflow.networks")
+@serializable
 class Time2Vec(keras.Layer):
     """
     Implements the Time2Vec learnbale embedding from [1].