<jupyter_start><jupyter_text>GPflow with TensorFlow 2=== Small steps big changes<jupyter_code>from typing import Tuple, Optional
from pathlib import Path

import datetime
import io
import matplotlib.pyplot as plt

import numpy as np
import tensorflow as tf
import gpflow

from gpflow.config import default_float

import warnings

warnings.filterwarnings('ignore')<jupyter_output><empty_output><jupyter_text>Make `tensorboard` work inside notebook:<jupyter_code>output_logdir = "/tmp/tensorboard"

!rm -rf "{output_logdir}"
!mkdir "{output_logdir}"

%load_ext tensorboard
%matplotlib inline


def enumerated_logdir(_logdir_id: int = [0]):
    logdir = Path(output_logdir, str(_logdir_id[0]))
    _logdir_id[0] += 1
    return str(logdir)<jupyter_output><empty_output><jupyter_text>Set up random seeds and default float for `gpflow` tensors:<jupyter_code>gpflow.config.set_default_float(np.float64)
np.random.seed(0)
tf.random.set_seed(0)<jupyter_output><empty_output><jupyter_text>Loading data using TensorFlow DatasetsFor this example, we create a synthetic dataset (noisy sine function):<jupyter_code>def noisy_sin(x): 
    return tf.math.sin(x) + 0.1 * tf.random.normal(x.shape, dtype=default_float())

num_train_data, num_test_data = 100, 500

X = tf.random.uniform((num_train_data, 1), dtype=default_float()) * 10
Xtest = tf.random.uniform((num_test_data, 1), dtype=default_float()) * 10

Y = noisy_sin(X)
Ytest = noisy_sin(Xtest)

data = (X, Y)

plt.plot(X, Y, 'xk')
plt.show()<jupyter_output><empty_output><jupyter_text>Working with TensorFlow Datasets is an efficient way to rapidly shuffle, iterate, and batch from data.<jupyter_code>train_dataset = tf.data.Dataset.from_tensor_slices((X, Y))
test_dataset = tf.data.Dataset.from_tensor_slices((Xtest, Ytest))

batch_size = 32
num_features = 10
prefetch_size = num_train_data // 2
shuffle_buffer_size = num_train_data // 2
num_batches_per_epoch = num_train_data // batch_size

original_train_dataset = train_dataset
train_dataset = train_dataset.repeat()\
                    .prefetch(prefetch_size)\
                    .shuffle(buffer_size=shuffle_buffer_size)\
                    .batch(batch_size)

print(f"prefetch_size={prefetch_size}")
print(f"shuffle_buffer_size={shuffle_buffer_size}")
print(f"num_batches_per_epoch={num_batches_per_epoch}")<jupyter_output>prefetch_size=50
shuffle_buffer_size=50
num_batches_per_epoch=3<jupyter_text>Define a GP modelIn GPflow 2.0, we use `tf.Module` (or the very thin `gpflow.base.Module` wrapper) to build all our models, as well as their components (kernels, likelihoods, parameters, and so on).<jupyter_code>kernel = gpflow.kernels.SquaredExponential(variance=2.)
likelihood = gpflow.likelihoods.Gaussian()
inducing_variable = np.linspace(0, 10, num_features).reshape(-1, 1)

model = gpflow.models.SVGP(kernel=kernel, likelihood=likelihood, inducing_variable=inducing_variable)<jupyter_output><empty_output><jupyter_text>You can set a module (or a particular parameter) to be non-trainable using the auxiliary method ```set_trainable(module, False)```:<jupyter_code>from gpflow.utilities import set_trainable

set_trainable(likelihood, False)
set_trainable(kernel.variance, False)

set_trainable(likelihood, True)
set_trainable(kernel.variance, True)<jupyter_output><empty_output><jupyter_text>We can use ```param.assign(value)``` to assign a value to a parameter:<jupyter_code>kernel.lengthscale.assign(0.5)<jupyter_output><empty_output><jupyter_text>All these changes are reflected when we use ```print_summary(model)``` to print a detailed summary of the model. By default the output is displayed in a minimalistic and simple table.<jupyter_code>from gpflow.utilities import print_summary

print_summary(model)  # same as print_summary(model, fmt="simple")<jupyter_output>name                      class      transform       trainable    shape        dtype    value
------------------------  ---------  --------------  -----------  -----------  -------  ----------------
SVGP.kernel.variance      Parameter  Softplus        True         ()           float64  2.0
SVGP.kernel.lengthscale   Parameter  Softplus        True         ()           float64  0.5
SVGP.likelihood.variance  Parameter  Softplus        True         ()           float64  1.0
SVGP.inducing_variable.Z  Parameter                  True         (10, 1)      float64  [[0....
SVGP.q_mu                 Parameter                  True         (10, 1)      float64  [[0....
SVGP.q_sqrt               Parameter  FillTriangular  True         (1, 10, 10)  float64  [[[1., 0., 0....<jupyter_text>We can change default printing so that it will look nicer in our notebook:<jupyter_code>gpflow.config.set_default_summary_fmt("notebook")

print_summary(model)  # same as print_summary(model, fmt="notebook")<jupyter_output><empty_output><jupyter_text>Training using Gradient TapesIn TensorFlow 2, we can optimize (trainable) model parameters with TensorFlow optimizers using `tf.GradientTape`. In this simple example, we perform one gradient update of the Adam optimizer to minimize the negative marginal log likelihood (or ELBO) of our model.<jupyter_code>optimizer = tf.optimizers.Adam()

with tf.GradientTape() as tape:
    tape.watch(model.trainable_variables)
    obj = - model.elbo(data)
    grads = tape.gradient(obj, model.trainable_variables)
    
optimizer.apply_gradients(zip(grads, model.trainable_variables))<jupyter_output><empty_output><jupyter_text>For a more elaborate example of a gradient update we can define an ```optimization_step``` that uses the decorator ```tf.function``` on a closure. A closure is a callable that returns the model objective evaluated at a given dataset when called.<jupyter_code>def optimization_step(model: gpflow.models.SVGP, batch: Tuple[tf.Tensor, tf.Tensor]):
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(model.trainable_variables)
        obj = - model.elbo(batch)
        grads = tape.gradient(obj, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))<jupyter_output><empty_output><jupyter_text>We can use the functionality of TensorFlow Datasets to define a simple training loop that iterates over batches of the training dataset:<jupyter_code>def simple_training_loop(model: gpflow.models.SVGP, epochs: int = 1, logging_epoch_freq: int = 10):
    batches = iter(train_dataset)
    tf_optimization_step = tf.function(optimization_step, autograph=False)
    for epoch in range(epochs):
        for _ in range(num_batches_per_epoch):
            tf_optimization_step(model, next(batches))

        epoch_id = epoch + 1
        if epoch_id % logging_epoch_freq == 0:
            tf.print(f"Epoch {epoch_id}: ELBO (train) {model.elbo(data)}")
simple_training_loop(model, epochs=10, logging_epoch_freq=2)<jupyter_output>Epoch 2: ELBO (train) -210.86251192675192
Epoch 4: ELBO (train) -209.20892625785794
Epoch 6: ELBO (train) -207.5394332593566
Epoch 8: ELBO (train) -205.86180178943908
Epoch 10: ELBO (train) -204.17336105880824<jupyter_text>Monitoring We can monitor the training procedure using `tf.summary`. First we create a summary writer object through which we can write scalars and images.<jupyter_code>from intro_to_gpflow2_plotting import plotting_regression, summary_matplotlib_image

samples_input = tf.cast(np.linspace(0, 10, 100).reshape(100, 1), default_float())

def monitored_training_loop(model: gpflow.models.SVGP, logdir: str, 
                            epochs: int = 1, logging_epoch_freq: int = 10,
                            num_samples: int = 10):
    summary_writer = tf.summary.create_file_writer(logdir)
    tf_optimization_step = tf.function(optimization_step)
    batches = iter(train_dataset)

    with summary_writer.as_default():
        for epoch in range(epochs):
            for _ in range(num_batches_per_epoch):
                tf_optimization_step(model, next(batches))

            epoch_id = epoch + 1
            if epoch_id % logging_epoch_freq == 0:
                tf.print(f"Epoch {epoch_id}: ELBO (train) {model.elbo(data)}")

                mean, var = model.predict_f(samples_input)
                samples = model.predict_f_samples(samples_input, num_samples)
                fig = plotting_regression(X, Y, samples_input, mean, var, samples)
                
                summary_matplotlib_image(dict(model_samples=fig), step=epoch)
                tf.summary.scalar('elbo', data=model.elbo(data), step=epoch)
                tf.summary.scalar('likelihood/variance', data=model.likelihood.variance, step=epoch)
                tf.summary.scalar('kernel/lengthscale', data=model.kernel.lengthscale, step=epoch)
                tf.summary.scalar('kernel/variance', data=model.kernel.variance, step=epoch)
model = gpflow.models.SVGP(kernel=kernel, likelihood=likelihood, inducing_variable=inducing_variable)

output_logdir = enumerated_logdir()
monitored_training_loop(model, output_logdir, epochs=1000, logging_epoch_freq=100)<jupyter_output>Epoch 100: ELBO (train) -141.1766391586027
Epoch 200: ELBO (train) -115.73096146310809
Epoch 300: ELBO (train) -99.34844151546481
Epoch 400: ELBO (train) -84.81535664328624
Epoch 500: ELBO (train) -70.54100940239233
Epoch 600: ELBO (train) -56.266306616978326
Epoch 700: ELBO (train) -42.13861564627568
Epoch 800: ELBO (train) -28.375123481848224
Epoch 900: ELBO (train) -15.13797813882487
Epoch 1000: ELBO (train) -2.604648948302163<jupyter_text>Then, we can use TensorBoard to examine the training procedure in more detail<jupyter_code># %tensorboard --logdir "{output_logdir}"<jupyter_output><empty_output><jupyter_text>Checkpointing: saving and loading modelsWith the help of `tf.train.CheckpointManager` and `tf.train.Checkpoint`, we can checkpoint the model throughout the training procedure. Let's start with a simple example using checkpointing to save and load a `tf.Variable`:<jupyter_code>initial_value = 1.2
a = tf.Variable(initial_value)<jupyter_output><empty_output><jupyter_text>Create `Checkpoint` object:<jupyter_code>ckpt = tf.train.Checkpoint(a=a)
manager = tf.train.CheckpointManager(ckpt, output_logdir, max_to_keep=3)<jupyter_output><empty_output><jupyter_text>Save the variable `a` and change its value right after:<jupyter_code>manager.save()
_ = a.assign(0.33)<jupyter_output><empty_output><jupyter_text>Now we can restore the old variable value:<jupyter_code>print(f"Current value of variable a: {a.numpy():0.3f}")

ckpt.restore(manager.latest_checkpoint)

print(f"Value of variable a after restore: {a.numpy():0.3f}")<jupyter_output>Current value of variable a: 0.330
Value of variable a after restore: 1.200<jupyter_text>In the example below, we modify a simple training loop to save the model every 100 epochs using the `CheckpointManager`.<jupyter_code>model = gpflow.models.SVGP(kernel=kernel, likelihood=likelihood, inducing_variable=inducing_variable)

def checkpointing_training_loop(model: gpflow.models.SVGP,
                                batch_size: int,
                                epochs: int,
                                manager: tf.train.CheckpointManager,
                                logging_epoch_freq: int = 100,
                                epoch_var: Optional[tf.Variable] = None,
                                step_var: Optional[tf.Variable] = None):
    tf_optimization_step = tf.function(optimization_step)
    batches = iter(train_dataset)
    
    for epoch in range(epochs):
        for step in range(num_batches_per_epoch):
            tf_optimization_step(model, next(batches))
            if step_var is not None:
                step_var.assign(epoch * num_batches_per_epoch + step + 1)
        if epoch_var is not None:
            epoch_var.assign(epoch + 1)
            
        epoch_id = epoch + 1
        if epoch_id % logging_epoch_freq == 0:
            ckpt_path = manager.save()
            tf.print(f"Epoch {epoch_id}: ELBO (train) {model.elbo(data)}, saved at {ckpt_path}")
step_var = tf.Variable(1, dtype=tf.int32, trainable=False)
epoch_var = tf.Variable(1, dtype=tf.int32, trainable=False)
ckpt = tf.train.Checkpoint(model=model, step=step_var, epoch=epoch_var)
manager = tf.train.CheckpointManager(ckpt, output_logdir, max_to_keep=5)

print(f"Checkpoint folder path at: {output_logdir}")

checkpointing_training_loop(model, batch_size=batch_size, epochs=1000, manager=manager, epoch_var=epoch_var, step_var=step_var)<jupyter_output>Checkpoint folder path at: /tmp/tensorboard/0
Epoch 100: ELBO (train) -39.88552537271265, saved at /tmp/tensorboard/0/ckpt-1
Epoch 200: ELBO (train) -18.616748229948485, saved at /tmp/tensorboard/0/ckpt-2
Epoch 300: ELBO (train) -9.922687181792332, saved at /tmp/tensorboard/0/ckpt-3
Epoch 400: ELBO (train) -2.8622609405418054, saved at /tmp/tensorboard/0/ckpt-4
Epoch 500: ELBO (train) 4.135491987946555, saved at /tmp/tensorboard/0/ckpt-5
Epoch 600: ELBO (train) 11.35093376190215, saved at /tmp/tensorboard/0/ckpt-6
Epoch 700: ELBO (train) 18.69857023728744, saved at /tmp/tensorboard/0/ckpt-7
Epoch 800: ELBO (train) 26.009422912807832, saved at /tmp/tensorboard/0/ckpt-8
Epoch 900: ELBO (train) 33.02582505540525, saved at /tmp/tensorboard/0/ckpt-9
Epoch 1000: ELBO (train) 39.49885808024786, saved at /tmp/tensorboard/0/ckpt-10<jupyter_text>After the models have been saved, we can restore them using ```tf.train.Checkpoint.restore``` and assert that their performance corresponds to that logged during training.<jupyter_code>for i, recorded_checkpoint in enumerate(manager.checkpoints):
    ckpt.restore(recorded_checkpoint)
    print(f"{i} restored model from epoch {int(epoch_var)} [step:{int(step_var)}] : ELBO training set {model.elbo(data)}")<jupyter_output>0 restored model from epoch 600 [step:1800] : ELBO training set 11.35093376190215
1 restored model from epoch 700 [step:2100] : ELBO training set 18.69857023728744
2 restored model from epoch 800 [step:2400] : ELBO training set 26.009422912807832
3 restored model from epoch 900 [step:2700] : ELBO training set 33.02582505540525
4 restored model from epoch 1000 [step:3000] : ELBO training set 39.49885808024786
{'fraction_non_alphanumeric': 0.1052407630664492, 'fraction_numerical': 0.05481880717679597, 'mean_word_length': 4.186465612357484, 'pattern_counts': {'":': 0, '<': 78, '<?xml version=': 0, '>': 78, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '21', 'dir': 'jupyter-structured-clean-dedup', 'id': '309594', 'license': 'Apache-2.0', 'n_tokens_mistral': 5135, 'n_tokens_neox': 4604, 'n_words': 1093, 'path': 'doc/source/notebooks/intro_to_gpflow2.ipynb', 'repo_name': 'elvijs/GPflow', 'size': '14101'}