替换tensorflow 2.0中损失函数的梯度计算答案

【问题标题】：Replacing gradient calculation of loss function in tensorflow 2.0替换tensorflow 2.0中损失函数的梯度计算
【发布时间】：2020-03-16 03:11:18
【问题描述】：

我想将tensorflow 2.0 中的损失函数替换为梯度函数。

例如，我有一个损失函数，如下所示：

def loss_function(prediction):
    # do some standard tensorflow things here
    return loss

然后我使用tf.GradientTape 方法应用渐变，即

with tf.GradientTape() as tape:

     prediction = model(input)
     loss = loss_function(prediction)

gradients = tf.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))

我的问题是我想更改梯度计算并自己显式计算它，仅针对当前自动计算的loss_function()。

我猜这与 @tf.custom_gradient 装饰器有关，但不确定如何让它为损失工作。

我正在使用与顺序/功能 api 相对的自定义训练循环。

【问题讨论】：

标签： tensorflow tensorflow2.0 loss-function

【解决方案1】：

答案其实很直接。首先定义一个 @tf.custom_gradient() 函数，它定义了梯度并通过它传递损失，即

@tf.custom_gradient
def custom_grad(x):
    def grad(dy, **kwargs):
        # do something with dy here
        return dy 
    return tf.identity(x), grad

并通过这个传递原始函数的损失变量：

def loss_function(prediction):
    # calculate the loss
    return custom_grad(loss)

【讨论】：