Tensorflow 2.0自定义梯度函数

我试图为我在论文中找到的分段连续函数建立一个自定义梯度

对数函数、线性函数和指数函数之间的连续体及其改善神经网络泛化的潜力-Godfrey

当前问题：在我的网络的自定义层中，可学习参数的梯度爆炸，导致生成

nans

。这个自定义层是在上面的文章中定义的一个可学习的激活函数

尝试过的解决方案

• 剪裁渐变优化器->当我使用例如

  keras.optimizer.Adam（clipnorm=0.5）

• 剪裁自定义层中learnable参数的值->这会产生一些我无法回忆的错误
• 创建自定义渐变计算以控制learnable参数的渐变->不确定应如何计算渐变

我现在的梯度计算函数是这样的

@tf.custom_gradient
def call_lt0(alpha_actv, x):
    def grad(dx):
        grad_x = 1/(1-alpha_actv * (alpha_actv + x))
        grad_alpha = (tf.math.log(1-alpha_actv * (alpha_actv + x)) - ((2 * alpha_actv ** 2 + alpha_actv * x)/(alpha_actv ** 2 + alpha_actv * x - 1))) / alpha_actv ** 2
        return grad_x * dx, grad_alpha * dx  # denk ik..?
    return (tf.math.exp(alpha_actv * x) - 1) / alpha_actv + alpha_actv, grad

这里我给出了x变量的梯度和可学习的参数alpha。但是，当我用

tf.GradientTape

检查这些输出时，我只得到一个输出。。因为我对

grad

的定义，我期望有两个。 另外，当我让keras/tensorflow计算所有内容，并再次使用

tf.GradientTape

检查输出时，我只得到一个值。。这就是我对单个可学习参数的期望。 所以在某个地方我做错了什么，或者误解了

grad

函数

---

有人能给我指路吗？

请提供模型和培训的代码示例。@thushv89我很想，但这可能会让我的文章有点臃肿。。我相信我的错误可能在于我的自定义梯度的定义或解释。例如我想学习

alpha_actv

参数，但要学习，我还需要

x

的值。在这个定义中，我还需要为

x

提供渐变，但这个渐变不需要任何更新。。那我是不是做错了什么？你的模型是什么样子的？完全连接？