tf.autodiff.ForwardAccumulator

Вычисляет произведения Якоби с вектором ("JVP") с помощью автоматической дифференциации в прямом направлении.

tf.autodiff.ForwardAccumulator(
    primals, tangents
)

Сравните с tf.GradientTape, которая вычисляет произведения вектор-Якоби ("VJP") с помощью автоматической дифференциации в обратном направлении (обратное распространение). Режим обратного распространения более привлекателен при вычислении градиентов скалярной функции по многим входным данным (например, нейронной сети со многими параметрами и скалярной функцией потерь). Режим прямого распространения лучше всего работает с функциями, имеющими много выходов и мало входов. Поскольку он не сохраняет промежуточные активации, он намного эффективнее с точки зрения памяти, чем обратное распространение, когда это применимо.

Рассмотрим простую линейную регрессию:

x = tf.constant([[2.0, 3.0], [1.0, 4.0]])
dense = tf.keras.layers.Dense(1)
dense.build([None, 2])
with tf.autodiff.ForwardAccumulator(
   primals=dense.kernel,
   tangents=tf.constant([[1.], [0.]])) as acc:
  loss = tf.reduce_sum((dense(x) - tf.constant([1., -1.])) ** 2.)
acc.jvp(loss)
<tf.Tensor: shape=(), dtype=float32, numpy=...>

В примере есть две переменные, содержащие параметры, dense.kernel (2 параметра) и dense.bias (1 параметр). Учитывая обучающие данные x как константу, это означает, что матрица Якоби для функции, отображающей параметры в функцию потерь, имеет одну строку и три столбца.

При прямом распространении мы заранее задаем вектор длиной 3, который умножает матрицу Якоби. Аргумент конструктора primals — это параметр (a tf.Tensor или tf.Variable), для которого мы задаем вектор, а аргумент tangents — это "вектор" в произведении Якоби с вектором. Если наша цель — вычислить всю матрицу Якоби, прямое распространение вычисляет один столбец за раз, в то время как обратное распространение вычисляет одну строку за раз. Поскольку матрица Якоби в примере линейной регрессии имеет только одну строку, обратное распространение требует меньше вызовов:

x = tf.constant([[2.0, 3.0], [1.0, 4.0]])
dense = tf.keras.layers.Dense(1)
dense.build([None, 2])
loss_fn = lambda: tf.reduce_sum((dense(x) - tf.constant([1., -1.])) ** 2.)
kernel_fprop = []
with tf.autodiff.ForwardAccumulator(
    dense.kernel, tf.constant([[1.], [0.]])) as acc:
  kernel_fprop.append(acc.jvp(loss_fn()))
with tf.autodiff.ForwardAccumulator(
    dense.kernel, tf.constant([[0.], [1.]])) as acc:
  kernel_fprop.append(acc.jvp(loss_fn()))
with tf.autodiff.ForwardAccumulator(dense.bias, tf.constant([1.])) as acc:
  bias_fprop = acc.jvp(loss_fn())
with tf.GradientTape() as tape:
  loss = loss_fn()
kernel_grad, bias_grad = tape.gradient(loss, (dense.kernel, dense.bias))
np.testing.assert_allclose(
    kernel_grad, tf.stack(kernel_fprop)[:, tf.newaxis])
np.testing.assert_allclose(bias_grad, bias_fprop[tf.newaxis])

Неявным в вызове tape.gradient является вектор длины 1, который умножает матрицу Якоби слева, т.е. произведение вектор-Якоби.

ForwardAccumulator сохраняет JVP, соответствующие примальным тензорам, за которыми он наблюдает, полученные из исходных primals, указанных в конструкторе. Как только примальный тензор удаляется, ForwardAccumulator удаляет соответствующий JVP.

acc.jvp(x) извлекает JVP acc, соответствующий примальному тензору x. Он не выполняет никаких вычислений. Вызовы acc.jvp могут повторяться, пока acc доступен, активен ли контекстный менеджер или нет. Новые JVP вычисляются только во время активности контекстного менеджера.

Обратите внимание, что JVP всегда применяются в том порядке, в котором их контекстные менеджеры были введены, поэтому внутренние аккумуляторы не увидят вычисления JVP от внешних аккумуляторов. Возьмите JVP высшего порядка из внешних аккумуляторов:

primal = tf.constant(1.1)
with tf.autodiff.ForwardAccumulator(primal, tf.constant(1.)) as outer:
  with tf.autodiff.ForwardAccumulator(primal, tf.constant(1.)) as inner:
    primal_out = primal ** tf.constant(3.5)
inner_jvp = inner.jvp(primal_out)
inner_jvp  # 3.5 * 1.1 ** 2.5
<tf.Tensor: shape=(), dtype=float32, numpy=4.4417057>
outer.jvp(inner_jvp)  # 3.5 * 2.5 * 1.1 ** 1.5
<tf.Tensor: shape=(), dtype=float32, numpy=10.094786>

Перестановка коллекции в последней строке для извлечения inner.jvp(outer.jvp(primal_out)) не сработает.

Строгое вложение также применяется к комбинациям ForwardAccumulator и tf.GradientTape. Более глубоко вложенные объекты GradientTape будут игнорировать продукты внешних объектов ForwardAccumulator. Это позволяет (например) эффективное с точки зрения памяти вычисление произведения Гессиана с вектором в прямом и обратном направлении, где внутренние объекты GradientTape иначе сохраняли бы все промежуточные JVP:

v = tf.Variable([1., 2.])
with tf.autodiff.ForwardAccumulator(
    v,
    # The "vector" in Hessian-vector product.
    tf.constant([1., 0.])) as acc:
  with tf.GradientTape() as tape:
    y = tf.reduce_sum(v ** 3.)
  backward = tape.gradient(y, v)
backward  # gradient from backprop
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([ 3., 12.], dtype=float32)>
acc.jvp(backward)  # forward-over-backward Hessian-vector product
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([6., 0.], dtype=float32)>

Методы

jvp

Просмотреть исходный код

jvp(
    primals, unconnected_gradients=tf.UnconnectedGradients.NONE
)

Извлекает произведение Якоби с вектором, вычисленное для primals.

Обратите внимание, что этот метод не выполняет никаких вычислений и просто ищет JVP, который уже был вычислен (в отличие от обратного распространения с использованием tf.GradientTape, где вычисление происходит при вызове tape.gradient).

__enter__

Просмотреть исходный код

__enter__()

__exit__

Просмотреть исходный код

__exit__(
    typ, value, traceback
)