I would like to compute the second derivatives (diagonal hessian) for all the components of all my variables in TensorFlow 2.0. I would like to autograph such a function as well.
I have it working in eager-mode on Google Colab (with a little test as well):
%tensorflow_version 2.x # for colab
import tensorflow as tf
x = tf.Variable([[1.], [2.]])
z = tf.Variable([[3., 4.]])
with tf.GradientTape(persistent=True) as tape:
with tf.GradientTape() as tape2:
y = (z @ x)**2
grads = tape2.gradient(y, [x, z])
# We want references to each component of our Variables in-order
# This needs to be done in the gradient tape, otherwise, we can't take
# gradients w.r.t. each component individually
grads_list = [list(tf.reshape(grad, [-1])) for grad in grads]
second_derivatives_list = []
for grad_list, var in zip(grads_list, [x, z]):
# Gradient with respect to x returns has the same shape as x
# So, we get the component that corresponds with the gradient in grads_list
temp = tf.stack([
tf.reshape(tape.gradient(g, var), [-1])[i] for i, g in enumerate(grad_list)
])
second_derivatives_list.append(tf.reshape(temp, var.shape))
del tape
assert list(second_derivatives_list[0].numpy().transpose()[0]) == [18., 32.]
assert list(second_derivatives_list[1].numpy()[0]) == [2., 8.]
But this implementation is slow and I cannot get it to work with Autograph. Does anyone have a better way to do this? Is there a way to take the gradient of a single component of a tensor without getting a reference to it in the tape? Thanks in advance!
Here are answers for TF 1.x: How to compute all second derivatives (only the diagonal of the Hessian matrix) in Tensorflow?, Tensorflow: Compute Hessian matrix (only diagonal part) with respect to a high rank tensor.
