4

我想计算 TensorFlow 2.0 中所有变量的所有分量的二阶导数（对角粗麻布）。我也想为这样的功能签名。

我让它在 Google Colab 上以渴望模式工作（也进行了一些测试）：

%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.]

但是这个实现很慢，我无法让它与 Autograph 一起使用。有没有人有更好的方法来做到这一点？有没有办法在不参考磁带的情况下获取张量的单个分量的梯度？提前致谢！

以下是 TF 1.x 的答案：如何在 Tensorflow 中计算所有二阶导数（仅 Hessian 矩阵的对角线）？, Tensorflow：计算关于高秩张量的 Hessian 矩阵（仅对角线部分）。

4

0 回答 0