小编Dav*_*gge的帖子

我目前正在尝试为自定义 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”的处理很慢。 …