python - 在 Tensorflow 2 中乘以常数值的变量没有梯度

Question

我正在创建一个自定义层，将内核乘以一个常量掩码值。掩码是不可训练的，但内核是可训练的。但在我的实现中，内核的渐变消失了。关于如何使这项工作的任何建议？

class MaskedDense(Dense):
    #Dense layer defined here:
    #https://github.com/tensorflow/tensorflow/blob/v2.4.0/tensorflow/python/keras/layers/core.py#L1081-L1247
    def __init__(self, units, mask_weight, mask_bias,
                 *args, **kwargs):
        super(MaskedDense, self).__init__(units, *args, **kwargs)        
        self.mask_weight = mask_weight
        self.mask_bias = mask_bias
        
    def build(self, input_shape):
        super(MaskedDense, self).build(input_shape)
        self._underlying_kernel = self.kernel
        self.kernel_mask = tf.Variable(initial_value=self.mask_weight, trainable=False, name='kernelmask', dtype='float32')
        self._non_trainable_weights.append(self.kernel_mask)
        self.kernel = tf.math.multiply(self._underlying_kernel, self.kernel_mask)
        
        if self.use_bias:
            self._underlying_bias = self.bias
            self.bias_mask = tf.Variable(initial_value=self.mask_bias, trainable=False, name='biasmask', dtype='float32')
            self._non_trainable_weights.append(self.bias_mask)
            self.bias = tf.math.multiply(self._underlying_bias, self.bias_mask)
        else:
            self.bias = None

我在一个简单的模型中为每一层输入一个掩码数组，如下所示：

mask_model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  MaskedDense(128, activation='relu', mask_weight = mask[0], mask_bias = mask[1]),
  MaskedDense(64, activation='relu', mask_weight = mask[2], mask_bias = mask[3]),
  MaskedDense(10, mask_weight = mask[4], mask_bias = mask[5])
])

使用这个模型给了我所有None的渐变。对使用原始层构建的模型使用相同的代码可以tensorflow.keras.layers.Dense正常工作。

标签： pythontensorflowkerastensorflow2.0