1

我正在研究 tensorflow 1.15 中的一个例程，该例程评估不同向量的几个粗麻布向量产品

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)来保留磁带的资源，但无法弄清楚如何将其合并到我的功能中。

4

1 回答 1

0

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

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

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

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

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

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

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 倍）。

于 2020-06-19T13:43:33.960 回答