Repository: secondmind-labs/trieste

Issue:
tensorboard logging model parameters or training stats for GPflux models



Discussion:

Resolving patch:
diff --git a/trieste/models/gpflux/models.py b/trieste/models/gpflux/models.py
--- a/trieste/models/gpflux/models.py
+++ b/trieste/models/gpflux/models.py
@@ -14,16 +14,18 @@
 
 from __future__ import annotations
 
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 import dill
 import gpflow
 import tensorflow as tf
 from gpflow.inducing_variables import InducingPoints
+from gpflow.utilities.traversal import _merge_leaf_components, leaf_components
 from gpflux.layers import GPLayer, LatentVariableLayer
 from gpflux.models import DeepGP
 from tensorflow.python.keras.callbacks import Callback
 
+from ... import logging
 from ...data import Dataset
 from ...types import TensorType
 from ..interfaces import (
@@ -60,13 +62,12 @@ def __init__(
         """
         :param model: The underlying GPflux deep Gaussian process model. Passing in a named closure
             rather than a model can help when copying or serialising.
-        :param optimizer: The optimizer configuration for training the model. Defaults to
-            :class:`~trieste.models.optimizer.KerasOptimizer` wrapper with
-            :class:`~tf.optimizers.Adam` optimizer. The ``optimizer`` argument to the wrapper is
-            used when compiling the model and ``fit_args`` is a dictionary of arguments to be used
-            in the Keras ``fit`` method. Defaults to 400 epochs, batch size of 1000, and verbose 0.
-            A custom callback that reduces the optimizer learning rate is used as well. See
-            https://keras.io/api/models/model_training_apis/#fit-method for a list of possible
+        :param optimizer: The optimizer wrapper with necessary specifications for compiling and
+            training the model. Defaults to :class:`~trieste.models.optimizer.KerasOptimizer` with
+            :class:`~tf.optimizers.Adam` optimizer, mean squared error metric and a dictionary of
+            default arguments for the Keras `fit` method: 400 epochs, batch size of 1000, and
+            verbose 0. A custom callback that reduces the optimizer learning rate is used as well.
+            See https://keras.io/api/models/model_training_apis/#fit-method for a list of possible
             arguments.
         :param num_rff_features: The number of random Fourier features used to approximate the
             kernel when calling :meth:`trajectory_sampler`. We use a default of 1000 as it typically
@@ -124,6 +125,9 @@ def scheduler(epoch: int, lr: float) -> float:
                 "callbacks": [tf.keras.callbacks.LearningRateScheduler(scheduler)],
             }
 
+        if self.optimizer.metrics is None:
+            self.optimizer.metrics = ["mse"]
+
         self._model_gpflux = model
         # inputs and targets need to be redone with a float64 dtype to avoid setting the keras
         # backend to float64, this is likely to be fixed in GPflux, see issue:
@@ -139,7 +143,7 @@ def scheduler(epoch: int, lr: float) -> float:
             dtype=tf.float64,
         )
         self._model_keras = model.as_training_model()
-        self._model_keras.compile(self.optimizer.optimizer)
+        self._model_keras.compile(self.optimizer.optimizer, metrics=self.optimizer.metrics)
         self._absolute_epochs = 0
         self._continuous_optimisation = continuous_optimisation
 
@@ -354,3 +358,71 @@ def optimize(self, dataset: Dataset) -> None:
         # so that the next time we call `optimize` the starting learning rate would be different.
         # Therefore we make sure the learning rate is set back to its initial value.
         self.optimizer.optimizer.lr.assign(self.original_lr)
+
+    def log(self, dataset: Optional[Dataset] = None) -> None:
+        """
+        Log model training information at a given optimization step to the Tensorboard.
+        We log a few summary statistics of losses, layer KL divergences and metrics (as provided in
+        ``optimizer``): ``final`` value at the end of the training, ``diff`` value as a difference
+        between inital and final epoch. We also log epoch statistics, but as histograms, rather
+        than time series. We also log several training data based metrics, such as root mean square
+        error between predictions and observations and several others.
+
+        For custom logs user will need to subclass the model and overwrite this method.
+
+        :param dataset: Optional data that can be used to log additional data-based model summaries.
+        """
+        summary_writer = logging.get_tensorboard_writer()
+        if summary_writer:
+            with summary_writer.as_default(step=logging.get_step_number()):
+                logging.scalar("epochs/num_epochs", len(self.model_keras.history.epoch))
+                # kernel parameters
+                for idx, layer in enumerate(self.model_gpflux.f_layers):
+                    kernel_components = _merge_leaf_components(leaf_components(layer.kernel))
+                    for k, v in kernel_components.items():
+                        if v.trainable:
+                            if tf.rank(v) == 0:
+                                logging.scalar(f"layer[{idx}]/kernel.{k}", v)
+                            elif tf.rank(v) == 1:
+                                for i, vi in enumerate(v):
+                                    logging.scalar(f"layer[{idx}]/kernel.{k}[{i}]", vi)
+                # likelihood parameters
+                likelihood_components = _merge_leaf_components(
+                    leaf_components(self.model_gpflux.likelihood_layer.likelihood)
+                )
+                for k, v in likelihood_components.items():
+                    if v.trainable:
+                        logging.scalar(f"likelihood/{k}", v)
+
+                # losses and metrics
+                for k, v in self.model_keras.history.history.items():
+                    logging.histogram(f"{k}/epoch", lambda: v)
+                    logging.scalar(f"{k}/final", lambda: v[-1])
+                    logging.scalar(f"{k}/diff", lambda: v[0] - v[-1])
+                # training data based diagnostics
+                if dataset:
+                    predict = self.predict(dataset.query_points)
+                    # training accuracy
+                    diffs = dataset.observations - predict[0]
+                    z_residuals = diffs / tf.math.sqrt(predict[1])
+                    logging.histogram("accuracy/absolute_errors", tf.math.abs(diffs))
+                    logging.histogram("accuracy/z_residuals", z_residuals)
+                    logging.scalar(
+                        "accuracy/root_mean_square_error", tf.math.sqrt(tf.reduce_mean(diffs ** 2))
+                    )
+                    logging.scalar(
+                        "accuracy/mean_absolute_error", tf.reduce_mean(tf.math.abs(diffs))
+                    )
+                    logging.scalar("accuracy/z_residuals_std", tf.math.reduce_std(z_residuals))
+                    # training variance
+                    variance_error = predict[1] - diffs ** 2
+                    logging.histogram("variance/predict_variance", predict[1])
+                    logging.histogram("variance/variance_error", variance_error)
+                    logging.scalar("variance/predict_variance_mean", tf.reduce_mean(predict[1]))
+                    logging.scalar(
+                        "variance/root_mean_variance_error",
+                        tf.math.sqrt(tf.reduce_mean(variance_error ** 2)),
+                    )
+                    # data stats
+                    empirical_variance = tf.math.reduce_variance(dataset.observations)
+                    logging.scalar("variance/empirical", empirical_variance)


Test patch:
diff --git a/tests/unit/models/gpflux/test_models.py b/tests/unit/models/gpflux/test_models.py
--- a/tests/unit/models/gpflux/test_models.py
+++ b/tests/unit/models/gpflux/test_models.py
@@ -26,6 +26,7 @@
 
 import copy
 import operator
+import unittest.mock
 from functools import partial
 from typing import Callable
 
@@ -43,6 +44,7 @@
 from tests.util.models.gpflux.models import single_layer_dgp_model
 from tests.util.models.models import fnc_2sin_x_over_3, fnc_3x_plus_10
 from trieste.data import Dataset
+from trieste.logging import step_number, tensorboard_writer
 from trieste.models.config import create_model
 from trieste.models.gpflux import DeepGaussianProcess
 from trieste.models.optimizer import KerasOptimizer
@@ -436,3 +438,70 @@ def test_deepgp_deep_copies_optimization_history() -> None:
     assert history.keys() == expected_history.keys()
     for k, v in expected_history.items():
         assert history[k] == v
+
+
+@unittest.mock.patch("trieste.logging.tf.summary.histogram")
+@unittest.mock.patch("trieste.logging.tf.summary.scalar")
+@pytest.mark.parametrize("use_dataset", [False, True])
+def test_deepgp_log(
+    mocked_summary_scalar: unittest.mock.MagicMock,
+    mocked_summary_histogram: unittest.mock.MagicMock,
+    use_dataset: bool,
+) -> None:
+    x_observed = np.linspace(0, 100, 100).reshape((-1, 1))
+    y_observed = fnc_2sin_x_over_3(x_observed)
+    dataset = Dataset(x_observed, y_observed)
+
+    model = DeepGaussianProcess(
+        single_layer_dgp_model(x_observed),
+        KerasOptimizer(tf.optimizers.Adam(), {"batch_size": 200, "epochs": 3, "verbose": 0}),
+    )
+    model.optimize(dataset)
+
+    mocked_summary_writer = unittest.mock.MagicMock()
+    with tensorboard_writer(mocked_summary_writer):
+        with step_number(42):
+            if use_dataset:
+                model.log(dataset)
+            else:
+                model.log(None)
+
+    assert len(mocked_summary_writer.method_calls) == 1
+    assert mocked_summary_writer.method_calls[0][0] == "as_default"
+    assert mocked_summary_writer.method_calls[0][-1]["step"] == 42
+
+    loss_names = ["loss/diff", "loss/final", "prior_kl/diff", "prior_kl/final"]
+    metrics_names = ["mse/diff", "mse/final"]
+    accuracy_names = [
+        "accuracy/root_mean_square_error",
+        "accuracy/mean_absolute_error",
+        "accuracy/z_residuals_std",
+    ]
+    variance_stats = [
+        "variance/predict_variance_mean",
+        "variance/empirical",
+        "variance/root_mean_variance_error",
+    ]
+    par_names = ["kernel.variance", "kernel.lengthscales", "likelihood"]
+
+    names_scalars = ["epochs/num_epochs"] + loss_names + metrics_names + par_names
+    num_scalars = len(names_scalars)
+    names_histogram = ["loss/epoch", "mse/epoch", "prior_kl/epoch"]
+    num_histogram = len(names_histogram)
+    if use_dataset:
+        names_scalars = names_scalars + accuracy_names + variance_stats
+        num_scalars = num_scalars + len(accuracy_names) + len(variance_stats)
+        names_histogram = (
+            names_histogram
+            + ["accuracy/absolute_errors", "accuracy/z_residuals"]
+            + ["variance/predict_variance", "variance/variance_error"]
+        )
+        num_histogram += 4
+
+    assert mocked_summary_scalar.call_count == num_scalars
+    for i in range(len(mocked_summary_scalar.call_args_list)):
+        assert any([j in mocked_summary_scalar.call_args_list[i][0][0] for j in names_scalars])
+
+    assert mocked_summary_histogram.call_count == num_histogram
+    for i in range(len(mocked_summary_histogram.call_args_list)):
+        assert any([j in mocked_summary_histogram.call_args_list[i][0][0] for j in names_histogram])