take batch size into account when aggregating metrics

Author: vpratz

bayesflow/approximators/approximator.py

@@ -137,3 +137,26 @@ def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = N
             self.build_from_data(mock_data)
 
         return super().fit(dataset=dataset, **kwargs)
+
+    def _batch_size_from_data(self, data: any):
+        """Obtain the batch size from a batch of data.
+
+        To properly weight the metrics for batches of different sizes, the batch size of a given batch of data is
+        required. As the data structure differs between approximators, each approximator has to specify this method.
+
+        Parameters
+        ----------
+        data :
+            The data that are passed to `compute_metrics` as keyword arguments.
+
+        Returns
+        -------
+        batch_size : int
+            The batch size of the given data.
+        """
+        raise NotImplementedError(
+            "Correct calculation of the metrics requires obtaining the batch size from the supplied data "
+            "for proper weighting of metrics for batches with different sizes. Please implement the "
+            "_batch_size_from_data method for your approximator. For a given batch of data, it should "
+            "return the corresponding batch size."
+        )

bayesflow/approximators/backend_approximators/jax_approximator.py

@@ -55,7 +55,7 @@ def stateless_test_step(self, state: tuple, data: dict[str, any]) -> (dict[str,
         )
         metrics, non_trainable_variables, metrics_variables = aux
 
-        metrics_variables = self._update_metrics(loss, metrics_variables)
+        metrics_variables = self._update_metrics(loss, metrics_variables, self._batch_size_from_data(data))
 
         state = trainable_variables, non_trainable_variables, metrics_variables
         return metrics, state
@@ -74,7 +74,7 @@ def stateless_train_step(self, state: tuple, data: dict[str, any]) -> (dict[str,
             optimizer_variables, grads, trainable_variables
         )
 
-        metrics_variables = self._update_metrics(loss, metrics_variables)
+        metrics_variables = self._update_metrics(loss, metrics_variables, self._batch_size_from_data(data))
 
         state = trainable_variables, non_trainable_variables, optimizer_variables, metrics_variables
         return metrics, state
@@ -85,11 +85,11 @@ def test_step(self, *args, **kwargs):
     def train_step(self, *args, **kwargs):
         return self.stateless_train_step(*args, **kwargs)
 
-    def _update_metrics(self, loss: jax.Array, metrics_variables: any) -> any:
+    def _update_metrics(self, loss: jax.Array, metrics_variables: any, sample_weight: any = None) -> any:
         # update the loss progress bar, and possibly metrics variables along with it
         state_mapping = list(zip(self.metrics_variables, metrics_variables))
         with keras.StatelessScope(state_mapping) as scope:
-            self._loss_tracker.update_state(loss)
+            self._loss_tracker.update_state(loss, sample_weight=sample_weight)
 
         metrics_variables = [scope.get_current_value(v) for v in self.metrics_variables]
 

bayesflow/approximators/backend_approximators/numpy_approximator.py

@@ -13,17 +13,17 @@ def compute_metrics(self, *args, **kwargs) -> dict[str, np.ndarray]:
     def test_step(self, data: dict[str, any]) -> dict[str, np.ndarray]:
         kwargs = filter_kwargs(data | {"stage": "validation"}, self.compute_metrics)
         metrics = self.compute_metrics(**kwargs)
-        self._update_metrics(metrics)
+        self._update_metrics(metrics, self._batch_size_from_data(data))
         return metrics
 
     def train_step(self, data: dict[str, any]) -> dict[str, np.ndarray]:
         raise NotImplementedError("Numpy backend does not support training.")
 
-    def _update_metrics(self, metrics):
+    def _update_metrics(self, metrics, sample_weight=None):
         for name, value in metrics.items():
             try:
                 metric_index = self.metrics_names.index(name)
-                self.metrics[metric_index].update_state(value)
+                self.metrics[metric_index].update_state(value, sample_weight=sample_weight)
             except ValueError:
                 self._metrics.append(keras.metrics.Mean(name=name))
-                self._metrics[-1].update_state(value)
+                self._metrics[-1].update_state(value, sample_weight=sample_weight)

bayesflow/approximators/backend_approximators/tensorflow_approximator.py

@@ -13,7 +13,7 @@ def compute_metrics(self, *args, **kwargs) -> dict[str, tf.Tensor]:
     def test_step(self, data: dict[str, any]) -> dict[str, tf.Tensor]:
         kwargs = filter_kwargs(data | {"stage": "validation"}, self.compute_metrics)
         metrics = self.compute_metrics(**kwargs)
-        self._update_metrics(metrics)
+        self._update_metrics(metrics, self._batch_size_from_data(data))
         return metrics
 
     def train_step(self, data: dict[str, any]) -> dict[str, tf.Tensor]:
@@ -26,14 +26,14 @@ def train_step(self, data: dict[str, any]) -> dict[str, tf.Tensor]:
         grads = tape.gradient(loss, self.trainable_variables)
         self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
 
-        self._update_metrics(metrics)
+        self._update_metrics(metrics, self._batch_size_from_data(data))
         return metrics
 
-    def _update_metrics(self, metrics):
+    def _update_metrics(self, metrics, sample_weight=None):
         for name, value in metrics.items():
             try:
                 metric_index = self.metrics_names.index(name)
-                self.metrics[metric_index].update_state(value)
+                self.metrics[metric_index].update_state(value, sample_weight=sample_weight)
             except ValueError:
                 self._metrics.append(keras.metrics.Mean(name=name))
-                self._metrics[-1].update_state(value)
+                self._metrics[-1].update_state(value, sample_weight=sample_weight)

bayesflow/approximators/backend_approximators/torch_approximator.py

@@ -13,7 +13,7 @@ def compute_metrics(self, *args, **kwargs) -> dict[str, torch.Tensor]:
     def test_step(self, data: dict[str, any]) -> dict[str, torch.Tensor]:
         kwargs = filter_kwargs(data | {"stage": "validation"}, self.compute_metrics)
         metrics = self.compute_metrics(**kwargs)
-        self._update_metrics(metrics)
+        self._update_metrics(metrics, self._batch_size_from_data(data))
         return metrics
 
     def train_step(self, data: dict[str, any]) -> dict[str, torch.Tensor]:
@@ -34,14 +34,14 @@ def train_step(self, data: dict[str, any]) -> dict[str, torch.Tensor]:
         with torch.no_grad():
             self.optimizer.apply(gradients, trainable_weights)
 
-        self._update_metrics(metrics)
+        self._update_metrics(metrics, self._batch_size_from_data(data))
         return metrics
 
-    def _update_metrics(self, metrics):
+    def _update_metrics(self, metrics, sample_weight=None):
         for name, value in metrics.items():
             try:
                 metric_index = self.metrics_names.index(name)
-                self.metrics[metric_index].update_state(value)
+                self.metrics[metric_index].update_state(value, sample_weight=sample_weight)
             except ValueError:
                 self._metrics.append(keras.metrics.Mean(name=name))
-                self._metrics[-1].update_state(value)
+                self._metrics[-1].update_state(value, sample_weight=sample_weight)

bayesflow/approximators/continuous_approximator.py

@@ -491,3 +491,6 @@ def _log_prob(
             conditions=inference_conditions,
             **filter_kwargs(kwargs, self.inference_network.log_prob),
         )
+
+    def _batch_size_from_data(self, data: Mapping[str, any]):
+        return keras.ops.shape(data["inference_variables"])[0]

bayesflow/approximators/model_comparison_approximator.py

@@ -378,3 +378,6 @@ def summaries(self, data: Mapping[str, np.ndarray], **kwargs):
         summary_variables = keras.ops.convert_to_tensor(data_adapted["summary_variables"])
         summaries = self.summary_network(summary_variables, **filter_kwargs(kwargs, self.summary_network.call))
         return summaries
+
+    def _batch_size_from_data(self, data: Mapping[str, any]):
+        return keras.ops.shape(data["model_indices"])[0]