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

查看official documentation了解更多

在 Tensorflow 中计算雅可比矩阵

Computing jacobian matrix in Tensorflow

python

gradient

tensorflow