fit#

gpjax.fit.fit(*, model, objective, train_data, optim, key=jr.key(42), num_iters=100, batch_size=-1, log_rate=10, verbose=True, unroll=1, safe=True)[source]#

Train a Module model with respect to a supplied objective function. Optimisers used here should originate from Optax.

Example

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> import optax as ox
>>> import gpjax as gpx
>>>
>>> xtrain = jnp.linspace(0, 1, 50).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)
>>>
>>> meanf = gpx.mean_functions.Constant()
>>> kernel = gpx.kernels.RBF()
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
>>> posterior = prior * likelihood
>>>
>>> nmll = lambda p, d: -gpx.objectives.conjugate_mll(p, d)
>>> trained_model, history = gpx.fit(
...     model=posterior, objective=nmll, train_data=D,
...     optim=ox.adam(0.01), num_iters=100, verbose=False,
... )
Parameters:

  • model (Model) – The model Module to be optimised.

  • objective (Objective) – The objective function that we are optimising with respect to.

  • train_data (Dataset) – The training data to be used for the optimisation.

  • optim (GradientTransformation) – The Optax optimiser that is to be used for learning a parameter set.

  • num_iters (int) – The number of optimisation steps to run. Defaults to 100.

  • batch_size (int) – The size of the mini-batch to use. Defaults to -1 (i.e. full batch).

  • key (KeyArray) – The random key to use for the optimisation batch selection. Defaults to jr.key(42).

  • log_rate (int) – How frequently the objective function’s value should be printed. Defaults to 10.

  • verbose (bool) – Whether to print the training loading bar. Defaults to True.

  • unroll (int) – The number of unrolled steps to use for the optimisation. Defaults to 1.

  • safe (bool)

Returns:

A tuple comprising the optimised model and training history.

Return type:

tuple[Model, Array]