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

Обратите внимание, что 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
)