Python tf.autodiff.ForwardAccumulator用法及代码示例

使用 forward-mode autodiff 计算 Jacobian-vector 产品 ("JVP"s)。

用法

tf.autodiff.ForwardAccumulator(
    primals, tangents
)

与使用reverse-mode autodiff(反向传播)计算vector-Jacobian产品("VJP"s)的tf.GradientTape相比。在计算scalar-valued 函数相对于许多输入的梯度时(例如，具有许多参数和标量损失的神经网络)，反向模式更具吸引力。前向模式最适用于输出多而输入少的函数。由于它不保留中间激活，因此它比适用的反向传播更节省内存。

考虑一个简单的线性回归：

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 视为常数，这意味着从参数到损失的函数映射的雅可比矩阵有一行和三列。

使用 forwardprop，我们预先指定一个 length-three 向量，它乘以雅可比行列式。 primals 构造函数参数是我们为其指定向量的参数(tf.Tensor 或 tf.Variable)，而 tangents 参数是 Jacobian-vector 产品中的 "vector"。如果我们的目标是计算整个雅可比矩阵，则 forwardprop 一次计算一列，而 backprop 一次计算一行。由于线性回归示例中的雅可比矩阵只有一行，因此反向传播需要更少的调用：

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 调用中隐含的是一个length-one 向量，它是left-multiplies 雅可比行列式，一个vector-Jacobian 乘积。

ForwardAccumulator 维护 JVP 对应的原始张量，它是从构造函数中指定的原始 primals 派生的。一旦原始张量被删除，ForwardAccumulator 就会删除相应的 JVP。

acc.jvp(x) 检索与原始张量 x 对应的 acc 的 JVP。它不执行任何计算。只要 acc 可以访问，无论上下文管理器是否处于活动状态，都可以重复 acc.jvp 调用。仅当上下文管理器处于活动状态时才计算新的 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 对象的乘积。这允许(例如)Hessian-vector 产品的内存高效 forward-over-backward 计算，其中内部 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)>