TensorFlow 1.x自定义ReLU可训练参数无法更新问题求助

问题描述 #

在TensorFlow 1.x中实现自定义改进版ReLU，引入了两个可训练参数coeff和coeff1，但经过1000步训练后，这两个参数的值完全没有变化。但如果使用Prelu函数内注释的代码逻辑，参数则能正常更新。

原因分析 #

问题出在当前Prelu函数的核心计算逻辑上：

copied_variable = tf.math.maximum(copied_variable,0)/copied_variable

当copied_variable > 0时，计算结果为1；当copied_variable <= 0时，结果为0。这个操作的输出是离散的常量值，和coeff、coeff1之间没有连续的可导关系——无论coeff/coeff1如何变化，输出要么是0要么是1，对这两个参数的梯度始终为0。梯度消失后，优化器无法通过反向传播更新参数。

而注释里的代码逻辑，是让coeff和coeff1直接参与到tf.math.maximum之前的加权计算中，输出结果和参数有连续的依赖关系，梯度可以正常回传到参数，因此优化器能正常更新它们。

解决方案 #

修改自定义ReLU的计算逻辑，确保coeff和coeff1的变化能影响最终输出的梯度，保留参数与输出之间的连续可导关系。比如可以参考标准PReLU的思路，让参数参与到负区间的加权，或者调整当前的组合计算方式，避免输出变成与参数无关的常量。

修改后的Prelu函数示例 #

def Prelu(x, coeff, coeff1):
    s = int(x.shape[-1])
    # 拆分通道
    x1 = x[:,:,:,:s//2]
    x2 = x[:,:,:,s//2:]
    # 调整计算逻辑，让coeff/coeff1参与激活的加权，保留可导关系
    sop = tf.math.maximum(x1 * coeff + x2 * coeff1, 0)
    sop1 = tf.math.maximum(x1 * coeff - x2 * coeff1, 0)
    copied_variable = tf.concat([sop, sop1], axis=-1)
    return copied_variable

这个修改去掉了导致梯度消失的除法操作，让coeff和coeff1直接参与到激活值的计算中，输出结果随参数变化连续变化，梯度可以正常反向传播，优化器就能更新这两个参数了。

验证效果 #

修改后重新运行训练代码，会发现coeff和coeff1的值在训练后会发生变化，同时损失值也会正常下降。

内容的提问来源于stack exchange，提问作者psj