python - 重复使用 GradientTape 进行多个雅可比计算

Question

我正在尝试计算 TensorFlow 神经网络输出相对于其输入的雅可比行列式。使用 tf.GradientTape.jacobian 方法可以轻松实现这一点。 TensorFlow 文档中提供的简单示例如下:

with tf.GradientTape() as g:
  x  = tf.constant([1.0, 2.0])
  g.watch(x)
  y = x * x
jacobian = g.jacobian(y, x)

如果我只想计算输入张量 x 的单个实例的雅可比行列式，这很好。然而，我需要对 x 的各种实例重复评估这个雅可比行列式很多很多次。对于重要的雅可比计算(例如，具有非线性激活函数的深度卷积神经网络)，重复重新运行 GradientTape 计算并评估雅可比方法的成本非常高。我从TensorFlow documentation知道梯度(以及雅可比行列式)是通过自动微分计算的。我必须想象网络的分析梯度(通过自动微分计算)有一些内部存储，它是在给定的输入下进行评估的。

我的问题:我假设 TensorFlow 构建并存储(至少部分)计算雅可比行列式所需的分析梯度是否正确？如果是这样，有没有办法保存这个分析梯度并用新的输入重新评估雅可比行列式，而不必通过 GradientTape 方法重建它？

“持久”GradientTape 似乎无法解决此问题:它只允许针对计算的多个内部参数重复评估单个 GradientTape 实例。

Answer 1

也许您发现这有帮助:

我需要多次计算任意函数的雅可比矩阵。我的问题是我不恰本地使用了 GradientTape，但我发布的代码可能会帮助您或给您一些见解。我发布了一个使用基于 session 的 tf.gradient() 函数和现代 GriadientTape 方法计算雅可比的独立示例。在帮助下，我让它们在彼此相同的数量级内运行。

• 如果您的问题集中在尝试重用调用之间的中间计算以提高速度，那么我认为 Nick 的答案更适用。
• 如果您的问题集中在尝试使 GradientTape 与静态图一样快，那么请确保将其包装在 @tf.function 中，因为它就是这样做的。

查看我的问题:Abysmal tf.GradientTape performance compared to tf.gradients() for computing jacobians