How can I use minibatches with a non-variational GPR in gpflow?

Question

I have tried to adapt the instructions in this documentation to use minibatches for a training a GPR model, but nothing I have tried works. I cannot supply the batch iterator to the training_loss_closure method or use a batch iterator for model's data attribute. Is there a way to use minibatches with a non-variational model, like a GPR or SGPR, in gpflow?

Answer 1

You can construct the data tuple to be two tf.Variable objects (if you want to be able to have different-length minibatches, you can pass a shape=None or shape=(None, dim) argument). Something like

X = tf.Variable(np.zeros(0, input_dim), shape=tf.TensorShape(None, input_dim), dtype=gpflow.default_float())
Y = tf.Variable(np.zeros(0, output_dim), shape=tf.TensorShape(None, output_dim), dtype=gpflow.default_float())
model = gpflow.models.GPR((X, Y), kernel)

Then you can write a loss function that takes in the current batch, assigns it to the variables, and then returns the model loss, along the lines of

@tf.function
def loss(data_batch):
    model.data[0].assign(data_batch[0])
    model.data[1].assign(data_batch[1])
    return model.training_loss()

Note: While this is numerically doable, for the non-SVGP models this might not give you the correct answer (the gradient you compute from a batch might not be an unbiased estimate of the full-batch gradient).