Tensorflow自定义渐变不会反向传播

如何解决Tensorflow自定义渐变不会反向传播

我试图用tf.custom_gradient来理解张量流中的自定义梯度，所以我试图用其梯度重现一个简单的矩阵乘法，其中我只对权重矩阵的梯度感兴趣。

@tf.custom_gradient
def matrixmul(x,weight):
    res = tf.matmul(x,weight)

    def grad(dy):
        grad_x = None
        grad_w = tf.matmul(x,dy,transpose_a=True)
        return grad_x,grad_w

    return res,grad

只要在反向传播中不使用我的自定义渐变，这都会提供正确的渐变，但是当我进行多次乘法运算时，渐变不会传播回去，因此我得到None作为{{ 1}}。

例如：

grad_w1

with tf.GradientTape() as tape:
    temp = matrixmul(x,w1)
    pred = tf.matmul(temp,w2)
    loss = tf.reduce_sum(tf.norm(pred - y,axis=-1))
grad_w1 = tape.gradient(loss,w1)

第一个片段给出正确的渐变，但是第二个片段产生with tf.GradientTape() as tape: temp = tf.matmul(x,w1) pred = matrixmul(temp,w1) 。