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 как константу, это означает, что матрица Якоби для функции отображения параметров в функцию потерь имеет одну строку и три столбца.

При использовании метода прямого распространения мы заранее задаём вектор длины три, который умножает матрицу Якоби. Аргумент конструктора 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 неявно содержится вектор длины один, который слева умножает матрицу Якоби — это произведение вектор-Якобиан.

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
)