如何用TensorFlow custom_gradient定义ReLU？为何梯度结果不符？

这个问题我之前也碰到过，核心问题出在你对tf.custom_gradient的梯度逻辑理解，以及张量操作的细节上——咱们一步步拆解：

先看你第一个实现的问题 #

@tf.custom_gradient
def relu(x):
    def grad(dy):
        return tf.cond(tf.reshape(x, []) > 0, lambda: tf.cast(tf.reshape(1, dy.shape), tf.float32), lambda: tf.cast(tf.reshape(0, dy.shape), tf.float32))
    return tf.nn.relu(x), grad

这里有两个关键错误：

1. 错误地将张量压缩为标量判断：tf.reshape(x, [])会把任意形状的输入x变成一个标量，这意味着你用整个张量的“代表值”（其实是被扁平化后只剩一个元素）来决定所有位置的梯度。比如如果x是[1.0, -2.0]，reshape后变成标量1.0，那么整个梯度都会返回全1，完全忽略了第二个元素-2.0应该对应梯度0的情况。原生ReLU的梯度是逐元素独立判断的，每个元素的梯度只和自身是否大于0有关。
2. 生成全1/全0张量的方式错误：tf.reshape(1, dy.shape)并不是生成和dy同形状的全1张量，而是试图把标量1强行reshape成dy的形状，这在逻辑上就不对。正确的写法应该是tf.ones_like(dy)（生成和dy同形状同 dtype 的全1张量）或者tf.ones(dy.shape, dtype=dy.dtype)，全0同理用tf.zeros_like(dy)。

再看第二个实现的问题 #

@tf.custom_gradient
def relu(x):
    def grad(dy):
        return dy
    return tf.nn.relu(x), grad

这个实现的问题更直接：你直接把上游梯度dy原封不动返回，完全没有根据x的取值做任何门控。原生ReLU的梯度逻辑是：当x>0时，梯度等于上游梯度dy；当x<=0时，梯度为0。而你这里不管x是什么值，梯度都直接传递dy，相当于ReLU的激活操作对梯度没有任何影响，这和原生逻辑完全相悖。

正确的自定义ReLU梯度实现 #

要和原生tf.nn.relu的梯度行为一致，你需要实现逐元素的梯度掩码：

@tf.custom_gradient
def relu(x):
    def grad(dy):
        # 逐元素判断x是否大于0，生成布尔掩码后转为和dy同类型的数值，再和dy相乘
        return dy * tf.cast(x > 0, dy.dtype)
    return tf.nn.relu(x), grad

这个实现里，每个元素的梯度都会根据自身x的取值决定是否保留上游梯度，和原生ReLU的梯度逻辑完全一致。

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