Метод batch_jacobian

Метод batch_jacobian класса GradientTape вычисляет матрицу Якоби для каждого отдельного примера в пакете данных. В отличие от метода jacobian, который суммирует градиенты по всему пакету, данный метод сохраняет результат независимо для каждого элемента батча. Первым параметром передаётся целевой тензор, по которому строится якобиан. Вторым параметром указывается тензор или список тензоров, по которым ведётся дифференцирование. Третьим необязательным параметром можно задать режим параллельных вычислений parallel_iterations.

Синтаксис

GradientTape.batch_jacobian(target, sources, parallel_iterations=1)

Пример

Давайте вычислим матрицу Якоби для функции возведения в квадрат, применённой к каждому элементу пакета:

import tensorflow as tf

x = tf.constant([1.0, 2.0, 3.0])

with tf.GradientTape() as tape:
	tape.watch(x)
	y = x * x

res = tape.batch_jacobian(y, x)
print(res)

Результат выполнения кода:

tf.Tensor(
[[2. 0. 0.]
 [0. 4. 0.]
 [0. 6. 0.]], shape=(3, 3), dtype=float32)

Каждая строка полученной матрицы содержит частные производные соответствующего элемента пакета по всем входным переменным.

Пример

Давайте вычислим якобиан для двумерного входного тензора с несколькими переменными:

import tensorflow as tf

x = tf.constant([[1.0, 2.0], [3.0, 4.0]])

with tf.GradientTape() as tape:
	tape.watch(x)
	y = x * x

res = tape.batch_jacobian(y, x)
print(res)

Результат выполнения кода:

tf.Tensor(
[[[2. 0.]
  [0. 4.]]

 [[6. 0.]
  [0. 8.]]], shape=(2, 2, 2), dtype=float32)

Для каждого примера в пакете формируется отдельная матрица Якоби, что видно по размерности результата.

Смотрите также

класс GradientTape,
который записывает операции для автоматического дифференцирования
метод jacobian,
который вычисляет якобиан с суммированием по пакету
метод gradient,
который вычисляет градиент целевой функции
метод watch,
который начинает отслеживание тензора