Faster posterior predictions with a model loaded from disk

[ignaciotb]

Hi,

I've recently learned that I can obtain faster predictions with an SVGP using:

  F_mu, F_sigma = self.model.posterior().predict_f(X_test)
  mu, sigma = self.model.likelihood.predict_mean_and_var(F_mu, F_sigma)

Instead of:

  mu, sigma = self.model.predict_y(X_test)

However I'm not succeeding to obtain the same improvement in speed on a model saved to disk, when comparing:

    self.model.predict_y_compiled = tf.function(
        self.model.predict_y, input_signature=[
            tf.TensorSpec(shape=[None, 2], dtype=tf.float64)]
    )

against:

    self.model.predict_f_compiled = tf.function(
        self.model.posterior().predict_f, input_signature=[
            tf.TensorSpec(shape=[None, 2], dtype=tf.float64)]
    )

    self.model.likelihood.predict_mean_and_var_compiled = tf.function(
        self.model.likelihood.predict_mean_and_var, input_signature=[
            tf.TensorSpec(shape=[None, 1], dtype=tf.float64), 
            tf.TensorSpec(shape=[None, 1], dtype=tf.float64)]
    )

being used in an equivalent manner as in the original model, that is:

    model = tf.saved_model.load(path_outputs + '/svgp')
    f_mu, f_sigma = model.predict_f_compiled(X_test)
    mu, sigma = model.likelihood.predict_mean_and_var_compiled(f_mu, f_sigma)

In fact, when doing the above, the predictions are slower than with predict_y_compiled():

    mu, sigma = model.predict_y_compiled(X_test)

I'm not very experienced with saving models through tf. Is there a more correct way that would yield a faster prediction when loading my model?

Thanks

Update:
The function predict_f_compiled() actually is much faster, surprisingly it's the call to likelihood.predict_mean_and_var_compiled() that seems to take up most of the processing time.