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)>