python - 在Tensorflow中计算jacobian矩阵

我想用张量流计算雅可比矩阵。
我所拥有的：

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需要大量的时间。
有什么想法吗？

最佳答案

假设X和Y是张量流张量，Y依赖于X：

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

结果具有Y.shape + X.shape形状，并提供Y的每个元素相对于X的每个元素的偏导数。