Python tensorflow、自定义损耗、自定义梯度、温度变量导致错误

考虑实现一个需要临时变量实例化的自定义损失函数。如果我们需要实现自定义梯度，TF希望梯度函数有一个额外的输出，而梯度的分量应该和损失函数的输入一样多。也就是说，如果我的理解是正确的话。如有任何更正，我们将不胜感激

链接相关github问题，其中包含一个最小工作示例（MWE）和其他调试信息：

此处，从github帖子粘贴的MWE副本为：

import tensorflow as tf
# from custom_gradient import custom_gradient  # my corrected version
from tensorflow import custom_gradient


def layer(t, name):
    var = tf.Variable(1.0, dtype=tf.float32, use_resource=True, name=name)
    return t * var


@custom_gradient
def custom_gradient_layer(t):
    result = layer(t, name='outside')

    def grad(*grad_ys, variables=None):
        assert variables is not None
        print(variables)
        grads = tf.gradients(
            layer(t, name='inside'),
            [t, *variables],
            grad_ys=grad_ys,
        )
        grads = (grads[:1], grads[1:])
        return grads

    return result, grad

这将抛出

ValueError:没有足够的值来解压缩…

如果我的理解是正确的，通常对于伴随方法（反向模式autodiff），正向过程建立表达式树，反向过程我们计算梯度，梯度函数的值乘以函数的偏导数，我们将对其进行导数，这可能是一个复合函数。如果需要，我可以发布推荐信

所以，对于一个输入变量，我们有一个梯度的计算。在这里，TF期望2，即使我们只有一个输入变量，因为temp变量，这在某些情况下是不可避免的

我的MWE伪代码如下所示：

@tf.custom_gradient
def custom_loss(in):

    temp = tf.Variable(tf.zeros([2 * N - 1]), dtype = tf.float32)

    ## compute loss function
    ...

     def grad(df):
         grad = df * partial_derivative
         return grad

    return loss, grad

---

安德烈

我也有同样的问题。我发现添加trainable=False可以解决我的问题。例如，以下

import tensorflow as tf
@tf.custom_gradient
def custom_loss(x):

    temp = tf.Variable(1., dtype = tf.float32)

    loss = x*temp

     def grad(dL):
         grad = dL * temp
         return grad

    return loss, grad

Gets me error“TypeError:如果将@custom_gradient与使用变量的函数一起使用，则grad_fn必须接受关键字参数‘variables’。”

但如果我这样做，我就不会出错

import tensorflow as tf
@tf.custom_gradient
def custom_loss(x):

    temp = tf.Variable(1., dtype = tf.float32, trainable=False)

    loss = x*temp

     def grad(dL):
         grad = dL * temp
         return grad

    return loss, grad

希望这对你或其他人有帮助