TFP exact GP regression in eager mode

Question

I'm trying to perform exact GP regression using the TF2.0 eager mode, based on the original graph based example from https://colab.research.google.com/github/tensorflow/probability/blob/master/tensorflow_probability/examples/jupyter_notebooks/Gaussian_Process_Regression_In_TFP.ipynb

amplitude = (
    np.finfo(np.float64).tiny +
    tf.nn.softplus(tf.Variable(initial_value=1., name='amplitude', dtype=np.float64))
)
length_scale = (
    np.finfo(np.float64).tiny +
    tf.nn.softplus(tf.Variable(initial_value=1., name='length_scale', dtype=np.float64))
)
observation_noise_variance = (
    np.finfo(np.float64).tiny +
    tf.nn.softplus(tf.Variable(initial_value=1e-6,
                               name='observation_noise_variance',
                               dtype=np.float64))
)

kernel = tfk.ExponentiatedQuadratic(amplitude, length_scale)

gp = tfd.GaussianProcess(
    kernel=kernel,
    index_points=tf.expand_dims(x, 1),
    observation_noise_variance=observation_noise_variance
)

neg_log_likelihood = lambda: -gp.log_prob(y)

optimizer = tf.optimizers.Adam(learning_rate=.01)

num_iters = 1000
lls_ = np.zeros(num_iters, np.float64)
for i in range(num_iters):
    lls_[i] = neg_log_likelihood()
    optimizer.minimize(neg_log_likelihood, var_list=[amplitude, length_scale, observation_noise_variance])

However optimization fails with:

'tensorflow.python.framework.ops.EagerTensor' object has no attribute '_in_graph_mode'

And if I move the amplitude, length_scale and observation_noise_variance each to tf.Variable, like:

amplitude = tf.Variable(initial_value=1., name='amplitude', dtype=np.float64)
amplitude_ = (
    np.finfo(np.float64).tiny +
    tf.nn.softplus(amplitude)
)

Optimization fails with:

ValueError: No gradients provided for any variable: ['amplitude:0', 'length_scale:0', 'observation_noise_variance:0'].

What am I doing wrong?

score 0 · Answer 1 · answered Aug 15 '19 at 08:02

There is an issue with eager mode currently:

https://groups.google.com/a/tensorflow.org/d/msg/tfprobability/IlhL-fcv3yc/jpQc4ogcFwAJ

Workaround was to work with GradientTape explicitly:

amplitude_ = tf.Variable(initial_value=1., name='amplitude_', dtype=np.float64)
length_scale_ = tf.Variable(initial_value=1., name='length_scale_', dtype=np.float64)
observation_noise_variance_ = tf.Variable(initial_value=1e-6,
                                         name='observation_noise_variance_',
                                         dtype=np.float64)

@tf.function
def neg_log_likelihood():
    amplitude = np.finfo(np.float64).tiny + tf.nn.softplus(amplitude_)
    length_scale = np.finfo(np.float64).tiny + tf.nn.softplus(length_scale_)
    observation_noise_variance = np.finfo(np.float64).tiny + tf.nn.softplus(observation_noise_variance_)

    kernel = tfk.ExponentiatedQuadratic(amplitude, length_scale)

    gp = tfd.GaussianProcess(
        kernel=kernel,
        index_points=tf.expand_dims(x, 1),
        observation_noise_variance=observation_noise_variance
    )

    return -gp.log_prob(y)

optimizer = tf.optimizers.Adam(learning_rate=.01)

num_iters = 1000

nlls = np.zeros(num_iters, np.float64)
for i in range(num_iters):
    nlls[i] = neg_log_likelihood()
    with tf.GradientTape() as tape:
        loss = neg_log_likelihood()
    grads = tape.gradient(loss, [amplitude_, length_scale_, observation_noise_variance_])
    optimizer.apply_gradients(zip(grads, [amplitude_, length_scale_, observation_noise_variance_]))

TFP exact GP regression in eager mode

1 Answers1