Вычисление якобиевой матрицы в тензорном потоке

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.