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 — это параметр (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 вычисляются только во время активности менеджера контекста.

Обратите внимание, что ForwardAccumulator всегда применяются в порядке входа в соответствующие менеджеры контекста, поэтому внутренние аккумуляторы не будут получать результаты вычисления 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
)