fix: move additional metrics from approximator to networks

Supplying the additional metrics for inference and summary networks
via the approximators compile method caused problems during
deseralization (#497). This can be resolved nicely by moving the metrics
directly to the networks' constructors, analogous to how Keras normally
handles custom metrics in layers.

As summary networks and inference networks inherit from the respective
base classes, this change only requires minor adaptations.

Author: vpratz

bayesflow/approximators/continuous_approximator.py

@@ -1,13 +1,14 @@
 from collections.abc import Mapping, Sequence, Callable
 
 import numpy as np
+import warnings
 
 import keras
 
 from bayesflow.adapters import Adapter
 from bayesflow.networks import InferenceNetwork, SummaryNetwork
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, logging, split_arrays, squeeze_inner_estimates_dict
+from bayesflow.utils import filter_kwargs, split_arrays, squeeze_inner_estimates_dict
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 from .approximator import Approximator
@@ -97,18 +98,21 @@ def build_adapter(
     def compile(
         self,
         *args,
-        inference_metrics: Sequence[keras.Metric] = None,
-        summary_metrics: Sequence[keras.Metric] = None,
         **kwargs,
     ):
-        if inference_metrics:
-            self.inference_network._metrics = inference_metrics
+        if "inference_metrics" in kwargs:
+            warnings.warn(
+                "Supplying inference metrics to the approximator is no longer supported. "
+                "Please pass the metrics directly to the network using the metrics parameter.",
+                DeprecationWarning,
+            )
 
-        if summary_metrics:
-            if self.summary_network is None:
-                logging.warning("Ignoring summary metrics because there is no summary network.")
-            else:
-                self.summary_network._metrics = summary_metrics
+        if "summary_metrics" in kwargs:
+            warnings.warn(
+                "Supplying summary metrics to the approximator is no longer supported. "
+                "Please pass the metrics directly to the network using the metrics parameter.",
+                DeprecationWarning,
+            )
 
         return super().compile(*args, **kwargs)
 
@@ -227,16 +231,6 @@ def get_config(self):
 
         return base_config | serialize(config)
 
-    def get_compile_config(self):
-        base_config = super().get_compile_config() or {}
-
-        config = {
-            "inference_metrics": self.inference_network._metrics,
-            "summary_metrics": self.summary_network._metrics if self.summary_network is not None else None,
-        }
-
-        return base_config | serialize(config)
-
     def estimate(
         self,
         conditions: Mapping[str, np.ndarray],

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -10,7 +10,6 @@
     expand_right_as,
     find_network,
     jacobian_trace,
-    layer_kwargs,
     weighted_mean,
     integrate,
     integrate_stochastic,
@@ -141,7 +140,8 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
 
     def get_config(self):
         base_config = super().get_config()
-        base_config = layer_kwargs(base_config)
+        # base distribution is fixed and passed in constructor
+        base_config.pop("base_distribution")
 
         config = {
             "subnet": self.subnet,

bayesflow/networks/consistency_models/consistency_model.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, layer_kwargs, weighted_mean
+from bayesflow.utils import find_network, weighted_mean
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from ..inference_network import InferenceNetwork
@@ -109,7 +109,8 @@ def from_config(cls, config, custom_objects=None):
 
     def get_config(self):
         base_config = super().get_config()
-        base_config = layer_kwargs(base_config)
+        # base distribution is fixed and passed in constructor
+        base_config.pop("base_distribution")
 
         config = {
             "total_steps": self.total_steps,

bayesflow/networks/coupling_flow/coupling_flow.py

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

bayesflow/networks/flow_matching/flow_matching.py

@@ -9,7 +9,6 @@
     find_network,
     integrate,
     jacobian_trace,
-    layer_kwargs,
     optimal_transport,
     weighted_mean,
 )
@@ -138,7 +137,6 @@ def from_config(cls, config, custom_objects=None):
 
     def get_config(self):
         base_config = super().get_config()
-        base_config = layer_kwargs(base_config)
 
         config = {
             "subnet": self.subnet,

bayesflow/networks/inference_network.py

@@ -1,12 +1,16 @@
 import keras
+from collections.abc import Sequence
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import allow_batch_size
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 
+@serializable("bayesflow.networks")
 class InferenceNetwork(keras.Layer):
-    def __init__(self, base_distribution: str = "normal", **kwargs):
+    def __init__(self, base_distribution: str = "normal", *, metrics: Sequence[keras.Metric] = None, **kwargs):
+        self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
         self.base_distribution = find_distribution(base_distribution)
 
@@ -72,3 +76,13 @@ def compute_metrics(
                 metrics[metric.name] = metric(samples, x)
 
         return metrics
+
+    def get_config(self):
+        base_config = super().get_config()
+        base_config = layer_kwargs(base_config)
+        config = {"metrics": self.custom_metrics, "base_distribution": self.base_distribution}
+        return base_config | serialize(config)
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))

bayesflow/networks/summary_network.py

@@ -1,14 +1,17 @@
 import keras
+from collections.abc import Sequence
 
 from bayesflow.metrics.functional import maximum_mean_discrepancy
 from bayesflow.types import Tensor
 from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import sanitize_input_shape
-from bayesflow.utils.serialization import deserialize
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 
+@serializable("bayesflow.networks")
 class SummaryNetwork(keras.Layer):
-    def __init__(self, base_distribution: str = None, **kwargs):
+    def __init__(self, base_distribution: str = None, *, metrics: Sequence[keras.Metric] = None, **kwargs):
+        self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
         self.base_distribution = find_distribution(base_distribution)
 
@@ -17,7 +20,7 @@ def build(self, input_shape):
         x = keras.ops.zeros(input_shape)
         z = self.call(x)
 
-        if self.base_distribution is not None:
+        if self.base_distribution is not None and not self.base_distribution.built:
             self.base_distribution.build(keras.ops.shape(z))
 
     @sanitize_input_shape
@@ -51,6 +54,12 @@ def compute_metrics(self, x: Tensor, stage: str = "training", **kwargs) -> dict[
 
         return metrics
 
+    def get_config(self):
+        base_config = super().get_config()
+        base_config = layer_kwargs(base_config)
+        config = {"base_distribution": self.base_distribution, "metrics": self.custom_metrics}
+        return base_config | serialize(config)
+
     @classmethod
     def from_config(cls, config, custom_objects=None):
         return cls(**deserialize(config, custom_objects=custom_objects))

tests/test_networks/conftest.py

@@ -1,6 +1,7 @@
 import pytest
 
 from bayesflow.networks import MLP
+from bayesflow.metrics import RootMeanSquaredError
 
 
 @pytest.fixture()
@@ -12,6 +13,7 @@ def diffusion_model_edm_F():
         integrate_kwargs={"method": "rk45", "steps": 250},
         noise_schedule="edm",
         prediction_type="F",
+        metrics=[RootMeanSquaredError()],
     )
 
 
@@ -82,22 +84,32 @@ def flow_matching():
     return FlowMatching(
         subnet=MLP([8, 8]),
         integrate_kwargs={"method": "rk45", "steps": 100},
+        metrics=[RootMeanSquaredError()],
     )
 
 
 @pytest.fixture()
 def consistency_model():
     from bayesflow.networks import ConsistencyModel
 
-    return ConsistencyModel(total_steps=100, subnet=MLP([8, 8]))
+    return ConsistencyModel(
+        total_steps=100,
+        subnet=MLP([8, 8]),
+        metrics=[RootMeanSquaredError()],
+    )
 
 
 @pytest.fixture()
 def affine_coupling_flow():
     from bayesflow.networks import CouplingFlow
 
     return CouplingFlow(
-        depth=2, subnet="mlp", subnet_kwargs=dict(widths=[8, 8]), transform="affine", transform_kwargs=dict(clamp=1.8)
+        depth=2,
+        subnet="mlp",
+        subnet_kwargs=dict(widths=[8, 8]),
+        transform="affine",
+        transform_kwargs=dict(clamp=1.8),
+        metrics=[RootMeanSquaredError()],
     )
 
 
@@ -106,15 +118,24 @@ def spline_coupling_flow():
     from bayesflow.networks import CouplingFlow
 
     return CouplingFlow(
-        depth=2, subnet="mlp", subnet_kwargs=dict(widths=[8, 8]), transform="spline", transform_kwargs=dict(bins=8)
+        depth=2,
+        subnet="mlp",
+        subnet_kwargs=dict(widths=[8, 8]),
+        transform="spline",
+        transform_kwargs=dict(bins=8),
+        metrics=[RootMeanSquaredError()],
     )
 
 
 @pytest.fixture()
 def free_form_flow():
     from bayesflow.experimental import FreeFormFlow
 
-    return FreeFormFlow(encoder_subnet=MLP([16, 16]), decoder_subnet=MLP([16, 16]))
+    return FreeFormFlow(
+        encoder_subnet=MLP([16, 16]),
+        decoder_subnet=MLP([16, 16]),
+        metrics=[RootMeanSquaredError()],
+    )
 
 
 @pytest.fixture()
@@ -236,35 +257,35 @@ def generative_inference_network(request):
 def time_series_network(summary_dim):
     from bayesflow.networks import TimeSeriesNetwork
 
-    return TimeSeriesNetwork(summary_dim=summary_dim)
+    return TimeSeriesNetwork(summary_dim=summary_dim, metrics=[RootMeanSquaredError()])
 
 
 @pytest.fixture(scope="function")
 def time_series_transformer(summary_dim):
     from bayesflow.networks import TimeSeriesTransformer
 
-    return TimeSeriesTransformer(summary_dim=summary_dim)
+    return TimeSeriesTransformer(summary_dim=summary_dim, metrics=[RootMeanSquaredError()])
 
 
 @pytest.fixture(scope="function")
 def fusion_transformer(summary_dim):
     from bayesflow.networks import FusionTransformer
 
-    return FusionTransformer(summary_dim=summary_dim)
+    return FusionTransformer(summary_dim=summary_dim, metrics=[RootMeanSquaredError()])
 
 
 @pytest.fixture(scope="function")
 def set_transformer(summary_dim):
     from bayesflow.networks import SetTransformer
 
-    return SetTransformer(summary_dim=summary_dim)
+    return SetTransformer(summary_dim=summary_dim, metrics=[RootMeanSquaredError()])
 
 
 @pytest.fixture(scope="function")
 def deep_set(summary_dim):
     from bayesflow.networks import DeepSet
 
-    return DeepSet(summary_dim=summary_dim)
+    return DeepSet(summary_dim=summary_dim, metrics=[RootMeanSquaredError()])
 
 
 @pytest.fixture(

tests/test_two_moons/conftest.py

@@ -4,8 +4,9 @@
 @pytest.fixture()
 def inference_network():
     from bayesflow.networks import CouplingFlow
+    from bayesflow.metrics import MaximumMeanDiscrepancy
 
-    return CouplingFlow(depth=2, subnet="mlp", subnet_kwargs=dict(widths=(32, 32)))
+    return CouplingFlow(depth=2, subnet="mlp", subnet_kwargs=dict(widths=(32, 32)), metrics=[MaximumMeanDiscrepancy()])
 
 
 @pytest.fixture()

tests/test_two_moons/test_two_moons.py

@@ -13,13 +13,9 @@ def test_compile(approximator, random_samples, jit_compile):
 
 
 def test_fit(approximator, train_dataset, validation_dataset, batch_size):
-    from bayesflow.metrics import MaximumMeanDiscrepancy
     from bayesflow.networks import PointInferenceNetwork
 
-    inference_metrics = []
-    if not isinstance(approximator.inference_network, PointInferenceNetwork):
-        inference_metrics += [MaximumMeanDiscrepancy()]
-    approximator.compile(inference_metrics=inference_metrics)
+    approximator.compile()
 
     mock_data = train_dataset[0]
     mock_data = keras.tree.map_structure(keras.ops.convert_to_tensor, mock_data)