如何让TensorFlow Keras中ThresholdedReLU的theta成为可优化参数

我构建了一个seq-to-seq模型，用于输入N时段连续数据并预测M时段（M<N）连续响应数据。由于响应数据稀疏（通常为0），我希望将“较小的”预测值压缩为0，因此添加了ThresholdedReLU作为最后一层，模型代码如下：

# encoder
encoder_inputs_11 = tf.keras.layers.Input(shape=(n_past, n_features))
encoder_layr1_11 = tf.keras.layers.LSTM(encoderLayers[0], return_state=True)
encoder_outputs1_11 = encoder_layr1_11(encoder_inputs_11)
encoder_states1_11 = encoder_outputs1_11[1:]
# decoder
decoder_inputs_11 = tf.keras.layers.RepeatVector(n_future)(encoder_outputs1_11[0])
decoder_layr1_11 = tf.keras.layers.LSTM(decoderLayers[0], return_sequences=True)(decoder_inputs_11, initial_state = encoder_states1_11)
decoder_outputs1_11 = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(n_response))(decoder_layr1_11)
# threshold layer
thresh_11 = tf.keras.layers.ThresholdedReLU(theta=reluTheta)(decoder_outputs1_11)
# entire model
result = tf.keras.models.Model(encoder_inputs_11, thresh_11)

请问如何让tf.keras.layers.ThresholdedReLU(theta=reluTheta)中的theta成为可优化参数？

Keras原生的ThresholdedReLU层的theta是固定超参数，无法被训练。要实现可训练的阈值，需要自定义一个继承自tf.keras.layers.Layer的层，将theta定义为可训练变量：

1. 自定义可训练ThresholdedReLU层 #

class TrainableThresholdedReLU(tf.keras.layers.Layer):
    def __init__(self, initial_theta=1.0, **kwargs):
        super(TrainableThresholdedReLU, self).__init__(**kwargs)
        self.initial_theta = initial_theta

    def build(self, input_shape):
        # 定义可训练的theta参数，初始值设为传入的initial_theta
        self.theta = self.add_weight(
            name='theta',
            shape=(),  # 单阈值，若需为每个特征设独立阈值可改为(n_response,)
            initializer=tf.keras.initializers.Constant(self.initial_theta),
            trainable=True
        )
        super(TrainableThresholdedReLU, self).build(input_shape)

    def call(self, inputs):
        # 实现ThresholdedReLU逻辑：输入大于theta时保留原值，否则输出0
        return tf.where(inputs > self.theta, inputs, 0.0)

    def get_config(self):
        config = super(TrainableThresholdedReLU, self).get_config()
        config.update({'initial_theta': self.initial_theta})
        return config

2. 修改原模型代码 #

将原有的ThresholdedReLU替换为自定义层即可：

# encoder
encoder_inputs_11 = tf.keras.layers.Input(shape=(n_past, n_features))
encoder_layr1_11 = tf.keras.layers.LSTM(encoderLayers[0], return_state=True)
encoder_outputs1_11 = encoder_layr1_11(encoder_inputs_11)
encoder_states1_11 = encoder_outputs1_11[1:]
# decoder
decoder_inputs_11 = tf.keras.layers.RepeatVector(n_future)(encoder_outputs1_11[0])
decoder_layr1_11 = tf.keras.layers.LSTM(decoderLayers[0], return_sequences=True)(decoder_inputs_11, initial_state = encoder_states1_11)
decoder_outputs1_11 = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(n_response))(decoder_layr1_11)
# 替换为自定义可训练阈值层
thresh_11 = TrainableThresholdedReLU(initial_theta=reluTheta)(decoder_outputs1_11)
# entire model
result = tf.keras.models.Model(encoder_inputs_11, thresh_11)

说明 #

• 自定义层中theta默认是标量，若需要为每个响应特征设置独立阈值，可将shape参数改为(n_response,)。
• 训练时，theta会和模型其他参数一起参与梯度更新，自动适配数据的稀疏特性。

内容的提问来源于stack exchange，提问作者Dr. Andrew