应用资源时访问多个变量的梯度 [Tensorflow]

Question

我目前正在尝试为自定义 tensorflow 层实现自定义优化。没有详细说明，我添加了一个小代码示例，它说明了我当前的代码是如何工作的。重要的部分是calculate_gradients(variables, gradients, momentum)一个函数，它需要层中所有变量的变量值和梯度。此外，该计算包含在优化期间必须存储的中间结果。这解释了说明性momentum变量。这种行为对我来说使得使用@custom_gradient不可能，因为这不允许我将此中间结果传播到优化器，然后优化器必须将其返回到自定义梯度函数以用于计算下一组梯度。除非有人知道这将如何工作（问题一），否则我还没有找到解决方法。

model = build_model()
for data, target in data:
    with tf.GradientTape() as tap:
        gradients = tape.gradient(loss(model(data), target), model.trainable_variables)
    for layer in model.layers:
        layer_gradients = gradients[indices] # actual indexing is not important
        new_gradients = calculate_gradients(layer.variables, layer_gradients, momentum)
        for variable, grad in zip(layer.variables, new_gradients):
            variable.assign(grad)

尝试在 tensorflow 优化器中实现这一点，特别是通过替换_resource_apply_dense如文档 [1] 中所示，我在分层行为方面遇到了一些麻烦。特别是因为_resource_apply_dense需要一个变量和一个梯度。第二个代码片段说明了我想要做什么，但目前还没有找到做这种get_other_variables_and_gradients(var)行为的方法。此外，该解决方案将为每一层计算三次梯度，这是非常次优的。

def _resource_apply_dense(var, grad, apply_state):
    other_vars_and_grads = get_other_variables_and_gradients(var)
    calculate_gradients(zip((var, grad), other_vars_and_gards))
    var.assign(grad)

简而言之，我的第二个问题是：有没有人知道如何实现这种行为，甚至在没有冗余计算的情况下甚至更好的新方法的情况下做得更好。目前，当我在训练循环中做所有事情时，优化工作，如代码片段一所示。因此，这只是与 tensorflow 优化器范例和性能集成的一个案例，因为在大型 for 循环中对列表进行非常“pythony”的处理很慢。

[1] https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Optimizer

Answer 1

我已经回答了我自己的问题并正在回复以便存档。在重新审视这个问题并查看 github [1] 上的 tf.keras.optimizers.Optimizers 代码后，我发现了一个私有方法“_transform_unaggreated_gradients”。在应用梯度之前调用此方法并允许必要的行为。就我而言，这看起来像这样：

  def _transform_unaggregated_gradients(self, grads_and_vars):
     for layer in model.layers:
        layer_gradients, layer_variables = grads_and_Vars[indices] # actual indexing is not important
        new_gradients = calculate_gradients(layer.variables, layer_gradients, momentum)
        grads_and_vars[indices] = (new_gradients, layer_variables)
    return grads_and_vars

[1] https://github.com/tensorflow/tensorflow/blob/v2.5.0/tensorflow/python/keras/optimizer_v2/optimizer_v2.py#L115-L1414