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的单个实例的Jacobian，这是很好的。但是，对于x的各种实例，我需要多次、多次地对这个Jacobian进行评估。对于一个非平凡的雅可比计算(例如，对于具有非线性激活函数的深卷积神经网络)，重复重新运行GradientTape计算和评估jacobian方法是非常昂贵的。我从TensorFlow文档中得知，梯度(因此也是雅可比)是通过自动微分计算的。我不得不想象网络的分析梯度(由自动微分计算)的内部存储，它是在给定的输入时评估的。

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

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

Answer 1

也许你觉得这很有帮助：

我需要多次、多次地计算任意函数的jacobian。我的问题是我不恰当地使用了GradientTape，但是我发布的代码可能会帮助您或给您一些洞察力。我发布了一个使用基于会话的tf.gradient()函数和现代GriadientTape方法计算雅可比的自包含示例。在我的帮助下，我让他们在相同的数量级内奔跑。

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

见我的问题：与tf.gradients()相比，tf.gradients()在计算雅克比时性能差

Answer 2

我是否正确地假设TensorFlow构建并存储了计算Jacobian?所需的分析梯度(至少部分)

不-我想你一定是误解了什么关于自动区分的东西。

当tf中的每个初等运算都知道其输出相对于输入的解析导数时，当计算实际梯度或Jacobian值时，将附加值(输出的导数)的数值传递给反传运算，然后使用每个初等运算的解析公式和链规则计算更多的数值。

，如果是的话，有没有办法保存这个分析梯度，用新的输入重新评估Jacobian，而不必通过GradientTape方法重建它呢？

不是的。如果要计算新输入的梯度或雅可比值，则需要再次执行整个计算。对于深度神经网络来说，这是无法避免的。

顺便说一句，如果你取神经网络的损失函数相对于你网络的参数的梯度，计算梯度的时间是O(1)计算损失本身的成本。这是反向传播，也是反向模式自动微分美的一部分.但是如果你的网络有N个输出，并且你想要计算你的网络的全部雅可比，那将花费O(N)计算你网络输出的时间。这也许就是为什么计算雅可比的代价如此高昂的原因。

例如，如果你在MNIST上训练一个网络，而你的网络有10个输出，你把这些输出组合成一个损失函数，那么计算损失函数的梯度将花费O(1)时间，但是计算关于参数的10个输出的雅可比值将花费O( 10 )时间。