How can we use lbfgs_minimize in TensorFlow 2.0

user_10392393 0
Viewed 4263

I have been unable to reproduce this example from tensorflow having tensorflow 2.0 installed.

This is the original snippet:

# A high-dimensional quadratic bowl.
ndims = 60
minimum = np.ones([ndims], dtype='float64')
scales = np.arange(ndims, dtype='float64') + 1.0

# The objective function and the gradient.
def quadratic(x):
    value = tf.reduce_sum(scales * (x - minimum) ** 2)
    return value, tf.gradients(value, x)[0]

start = np.arange(ndims, 0, -1, dtype='float64')
optim_results = tfp.optimizer.lbfgs_minimize(
  quadratic, initial_position=start, num_correction_pairs=10,
  tolerance=1e-8)

with tf.Session() as session:
    results = session.run(optim_results)
# Check that the search converged
assert(results.converged)
# Check that the argmin is close to the actual value.
np.testing.assert_allclose(results.position, minimum)

Which does not work with the following error:

RuntimeError: tf.gradients is not supported when eager execution is enabled. Use tf.GradientTape instead.

If I change the code and have gradient tape instead like the following:

def quadratic(x):

    x = tf.Variable(x, dtype='float64')
    with tf.GradientTape() as t:
        value = tf.reduce_sum(scales * (x - minimum) ** 2)
    grad = t.gradient(value, x)

    return value, grad

I also get the following error:

TypeError: Tensor is unhashable if Tensor equality is enabled. Instead, use tensor.experimental_ref() as the key.

In general anything I tried didn't work and I don't know how I can use lbfgs in tensorflow 2.0.

tensorflow2.0 tensorflow deep-learning python
2 Answers

Use tf.function in your objective function so it is executed as a graph, then you will be able to use tf.gradients:

import tensorflow as tf
import tensorflow_probability as tfp
import numpy as np

# A high-dimensional quadratic bowl.
ndims = 60
minimum = tf.ones([ndims], dtype='float64')
scales = tf.range(ndims, dtype='float64') + 1.0

# The objective function and the gradient.
@tf.function
def quadratic(x):
    value = tf.reduce_sum(scales * (x - minimum) ** 2)
    return value, tf.gradients(value, x)[0]

start = tf.range(ndims, 0, -1, dtype='float64')
optim_results = tfp.optimizer.lbfgs_minimize(
  quadratic, initial_position=start, num_correction_pairs=10,
  tolerance=1e-8)

# Check that the search converged
print(optim_results.converged.numpy())
# True

# Check that the argmin is close to the actual value.
print(np.allclose(optim_results.position.numpy(), minimum.numpy()))
# True
user_17827920

Just to add a little to the answer by @jdehesa - it can also be useful to use tfp.math.value_and_gradient in this case, which will create the gradient tape for you if you are using eager mode. For example:

import tensorflow as tf
import tensorflow_probability as tfp

ndims = 60
minimum = tf.ones([ndims], dtype="float64")
scales = tf.range(ndims, dtype="float64") + 1.0


def quadratic(x):
    value = tf.reduce_sum(scales * (x - minimum) ** 2)
    return value


start = tf.range(ndims, 0, -1, dtype="float64")

optim_results = tfp.optimizer.lbfgs_minimize(
    lambda x: tfp.math.value_and_gradient(quadratic, x),
    initial_position=start,
    num_correction_pairs=10,
    tolerance=1e-8,
)
user_46155660
Related