在 Tensorflow 中計算雅可比矩陣

Question

我想通過 Tensorflow 計算雅可比矩陣。

我擁有的：

def compute_grads(fn, vars, data_num):
    grads = []
    for n in range(0, data_num):
        for v in vars:
            grads.append(tf.gradients(tf.slice(fn, [n, 0], [1, 1]), v)[0])
    return tf.reshape(tf.stack(grads), shape=[data_num, -1])

fn是損失函數， vars都是可訓練的變量， data_num是數據的數量。

但是如果我們增加數據的數量，運行函數compute_grads會花費大量時間。 有任何想法嗎？

Answer 1

假設X和Y是 Tensorflow 張量並且Y取決於X ：

from tensorflow.python.ops.parallel_for.gradients import jacobian
J=jacobian(Y,X)

結果具有形狀Y.shape + X.shape並提供Y的每個元素相對於X的每個元素的偏導數。

Answer 2

假設您使用的是 Tensorflow 2 或 Tensorflow <2 和 Eager 模式，您可以使用 GradientTape 和 inbuild 函數：

with tf.GradientTape() as g:
  x  = tf.constant([1.0, 2.0])
  g.watch(x)
  y = x * x
jacobian = g.jacobian(y, x)
# jacobian value is [[2., 0.], [0., 4.]]

查看官方文檔了解更多