如何在嵌套渐变磁带中重用内部渐变？

Question

我正在使用 tensorflow 1.15 中的一个例程来评估不同向量的几个 hessian-vector 产品

def hessian_v_prod(self, v):
    with tf.GradientTape() as t1:
        with tf.GradientTape() as t2:
            # evaluate loss which uses self.variables
            loss_val = self.loss()
        grad = t2.gradient(loss_val, self.variables)
        v_hat = tf.reduce(tf.multiply(v, grad))

    return t1.gradient(v_hat, self.variables)

每次调用此函数时，它都必须评估内部循环并计算梯度，但无论v 的值如何，这都是相同的。如何在每次调用此函数时重用 grad 值？

我看到有一个选项可以将磁带创建为 tf.GradientTape(persist=True)，它保留磁带的资源，但不知道如何将其合并到我的函数中。

Answer 1

我不得不深入研究GradientTape 的内部工作原理，但设法弄明白了。 在这里分享给其他可能有同样问题的人。 剧透警告：这有点骇人听闻！

首先，调用时实际发生了什么

with tf.GradientTape() as tape:
    loss_value = self.loss()
tape.gradient(loss_value, vars)

要找出这一点，我们需要检查分别在 with 块的开头和结尾调用的 __enter__() 和 __exit__() 函数。

在tensorflow_core/python/eager/backprop.py

def __enter__(self):
    """Enters a context inside which operations are recorded on this tape."""
    self._push_tape()
    return self

def __exit__(self, typ, value, traceback):
    """Exits the recording context, no further operations are traced."""
    if self._recording:
        self._pop_tape()

我们可以自己使用这些私有函数来控制录制，而不需要 with 块。

# Initialize outer and inner tapes
self.gt_outer = tf.GradientTape(persistent=True)
self.gt_inner = tf.GradientTape(persistent=True)

# Begin Recording
self.gt_outer._push_tape()
self.gt_inner._push_tape()

# evaluate loss which uses self.variables
loss_val = self.loss()

# stop recording on inner tape
self.gt_inner._pop_tape()

# Evaluate the gradient on the inner tape
self.gt_grad = self.gt_inner.gradient(loss_val, self.variables)

# Stop recording on the outer tape
self.gt_outer._pop_tape()

现在，每当我们需要评估 hessian 向量积时，我们都可以重复使用外部梯度带。

def hessian_v_prod(self, v):
    self.gt_outer._push_tape()
    v_hat = tf.reduce(tf.multiply(v, self.gt_grad))
    self.gt_outer._pop_tape()
    return self.gt_outer.gradient(v_hat, self.variables)

请注意，我们正在持久化磁带，因此每次计算 hessian 向量积时都会使用更多内存。无法保留部分磁带内存，因此在某些时候需要重置磁带。

# reset tapes
self.gt_outer._tape = None
self.gt_inner._tape = None

要在此之后再次使用它们，我们需要重新评估内部循环。它并不完美，但它可以完成这项工作并显着提高速度（接近 2 倍），但代价是更大的内存使用量。