Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Try using step_callback to get gpflow training loss history #772

Closed
wants to merge 1 commit into from
Closed

Try using step_callback to get gpflow training loss history #772

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions tests/unit/models/gpflow/test_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@
RandomSubSampleInducingPointSelector,
UniformInducingPointSelector,
)
from trieste.models.gpflow.optimizer import loss_history_callback
from trieste.models.gpflow.sampler import (
DecoupledTrajectorySampler,
RandomFourierFeatureTrajectorySampler,
Expand Down Expand Up @@ -155,6 +156,22 @@ def test_gpflow_wrappers_default_optimize(
assert internal_model.training_loss(**args) < loss


def test_gpflow_history_callback(
gpflow_interface_factory: ModelFactoryType,
) -> None:
data = mock_data()
optimizer = Optimizer(optimizer=gpflow.optimizers.Scipy())
model, _ = gpflow_interface_factory(*data, optimizer=optimizer)
losses: list[TensorType] = []
model.optimizer.minimize_args["step_callback"] = loss_history_callback(
model.model, data, losses
)
model.optimize(Dataset(*data))

assert len(losses) > 50
assert losses[-1] < losses[0]


def test_gpflow_wrappers_ref_optimize(gpflow_interface_factory: ModelFactoryType) -> None:
x = tf.constant(np.arange(5).reshape(-1, 1), dtype=gpflow.default_float())
y = fnc_2sin_x_over_3(x)
Expand Down
31 changes: 30 additions & 1 deletion trieste/models/gpflow/optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,14 @@

from __future__ import annotations

from typing import Any, Callable, Optional
from typing import Any, Callable, Optional, Sequence

import tensorflow as tf
from gpflow.models import ExternalDataTrainingLossMixin, InternalDataTrainingLossMixin
from gpflow.optimizers.scipy import StepCallback
from tensorflow.python.data.ops.iterator_ops import OwnedIterator as DatasetOwnedIterator

from ...types import TensorType
from ..optimizer import LossClosure, TrainingData, create_loss_function


Expand Down Expand Up @@ -92,3 +94,30 @@ def compiled_closure() -> tf.Tensor:
builder.closure_builder = closure_builder

return builder.closure_builder((X, Y))


def loss_history_callback(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking of getting the whole OptimizeResult object, rather than just loss: https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.OptimizeResult.html

when we call minimize in gpflow.scipy, it should return it, right?
then fun should be actual loss value, but one gets all the other stuff that is useful for evaluating the optimisation: success, status, nfev, nit ...

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can certainly get the OptimizeResult, but I thought that didn't include any loss history (unlike the keras optimizer) just the number of evaluations? Isn't that what you wanted?

Copy link
Collaborator Author

@uri-granta uri-granta Jul 24, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As an aside, how were you expecting to get the result? TrainableProbabilisticModel.optimize currently doesn't return anything (though Optimizer.optimize already does). If you're using BO or AskTell then there is no obvious way to return all the results when optimizing all the models, but I guess we could place them somewhere in the model wrappers (a bit like how the keras history is already available inside DeepEnsemble.model.history).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, you want to say we would get here a loss with every iteration of the optimiser? rather than a single end result with fun in OptimizeResult?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we could place them somewhere in the model wrappers (a bit like how the keras history is already available inside DeepEnsemble.model.history)

yes, this was the solution I was thinking of - but not sure if we can attach a history object to gpflow model like we have with keras model object, if not and we store it in the wrapper, perhaps we should have a unified way for all models, store it in the wrapper object

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could also just store the most recent result of calling Optimizer.optimize in the optimizer itself, which would let us easily access the OptimizeResult object. Or even all the previous results (enabled via a save_results: bool option in the Optimizer).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, you want to say we would get here a loss with every iteration of the optimiser? rather than a single end result with fun in OptimizeResult?

Yes, that's what this does. Essentially equivalent to history.history["loss"] on the keras result.

I see, that would be useful, but since it can expensive perhaps leaving off by default? (depends how expensive it is...)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could also just store the most recent result of calling Optimizer.optimize in the optimizer itself, which would let us easily access the OptimizeResult object. Or even all the previous results (enabled via a save_results: bool option in the Optimizer).

I think that is the best solution, all the history from each active learning step

for keras we can simply copy it from the model attribute, just so that we have a unified way of accessing the optimization history?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remember that for keas the history is generated by the model calling the fit method, not the optimizer's minimize method.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we could still attach the resulting history to the optimizer attribute, we would just need to do it from the model wrapper, in the optimize method, no?

model: InternalDataTrainingLossMixin | ExternalDataTrainingLossMixin,
data: TrainingData,
history: list[TensorType],
compile: bool = True,
) -> StepCallback:
"""
A step callback to use with the :class:`~gpflow.optimizers.Scipy` optimizer that tracks
the training loss history. To use this with a `GPflowPredictor model, assign the result
to `model.optimizer.minimize_args["step_callback"]`.

:param model: The model to track the loss function for.
:param dataset: The data with which to track the loss function.
:param history: A list to append loss values to.
:param compile: Whether to compile with :func:`tf.function`.
:return: A `StepCallback` to use with the :class:`~gpflow.optimizers.Scipy` optimizer.
"""
training_loss = create_loss_function(model, data, compile=compile)

def step_callback(
step: int, variables: Sequence[tf.Variable], values: Sequence[tf.Tensor]
) -> None:
history.append(training_loss())

return step_callback
Loading