在 Tensorflow 2.0 中使用自定义层时，渐变返回 None

Question

我设计了一个自定义层，像这样使用：

class SquaredWeightedLayer(tf.keras.layers.Layer):
    def __init__(self, units=1):
        super(SquaredWeightedLayer, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units), initializer="random_normal", trainable=True)
        self.square_w = self.w ** 2
        self.b = self.add_weight(shape=(self.units,), initializer="zeros", trainable=True)
        super(SquaredWeightedLayer, self).build(input_shape)

    def call(self, inputs):
        return tf.sigmoid(tf.add(tf.matmul(inputs, self.square_w), self.b))

但是，tape.gradient(loss, self.w) 返回 None 和 tape.gradient(loss, self.square_w) 返回正常值。 损失是binary_crossentropy。

我将非常感谢任何解决此问题的建议。谢谢！

Answer 1

问题是self.w ** 2的计算已经在build函数中执行，在任何磁带上下文之外，所以Tensorflow无法追踪square_w实际上来自w，导致@987654325的梯度@。您只需将平方操作移至call即可修复它：

def call(self, inputs):
    return tf.sigmoid(tf.add(tf.matmul(inputs, self.w**2), self.b))