tf.autodiff.ForwardAccumulator

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

tf.autodiff.ForwardAccumulator(
    primals, tangents
)

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

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

x = tf.constant([[2.0, 3.0], [1.0, 4.0]])
targets = tf.constant([[1.], [-1.]])
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) - targets) ** 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]])
targets = tf.constant([[1.], [-1.]])
dense = tf.keras.layers.Dense(1)
dense.build([None, 2])
loss_fn = lambda: tf.reduce_sum((dense(x) - targets) ** 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
)