Keras中如何实现依赖输入特征的自定义条件输出层激活函数

Keras默认的激活接口仅支持传入当前层的输出张量，你这类需要依赖原始输入特征的激活逻辑，不能用常规的activation参数传自定义函数实现，核心要注意两个要点：

• 放弃序贯（Sequential）模型，用函数式API搭建网络，才能跨层传递原始输入张量到输出层的计算节点
• 所有条件判断不能用Python原生的if/else，必须换成TensorFlow的张量算子实现向量化计算，否则图模式下会触发张量布尔值解析错误，也无法正常训练和导出模型。

具体实现步骤 #

1. 定义双特征输入层，单独拆分出x1、x2两个特征维度，注意保留维度对齐后续隐藏层输出，避免广播报错
2. 搭建输出层时不要预置激活，直接拿到输出层的原始输出张量y
3. 用tf.where重写激活逻辑，逐样本匹配条件分支计算结果
4. 将激活计算后的结果作为模型的最终输出，包裹成Model实例即可

可运行代码示例 #

import tensorflow as tf
from tensorflow.keras import layers, Model

# 定义输入层：输入为2个特征(x1, x2)
inputs = layers.Input(shape=(2,))
x1 = inputs[:, 0:1]  # 切片保留第二维度，形状为(batch_size, 1)，和隐藏层输出对齐
x2 = inputs[:, 1:2]

# 搭建隐藏层，隐藏层可正常使用内置激活
hidden = layers.Dense(32, activation='relu')(inputs)
# 输出层输出原始张量y，不设置activation参数
y = layers.Dense(1)(hidden)

# 实现自定义激活逻辑
# 逻辑：x1>1.2时输出0，否则输出min(y, x2)
condition = tf.greater(x1, 1.2)
output = tf.where(
    condition,
    tf.zeros_like(y),
    tf.math.minimum(y, x2)
)

# 组装模型
model = Model(inputs=inputs, outputs=output)

# 推理测试
test_input = tf.constant([[1.5, 2.0], [0.8, 3.0], [1.0, 0.5]])
print(model(test_input))

额外说明 #

• 如果实际激活逻辑非常复杂，无法直接用向量化算子拼接，可以自定义一个继承tf.keras.layers.Layer的激活层，在call方法中实现计算逻辑，调用时把y、x1、x2作为参数传入即可，封装性更好。
• 不要尝试用Python原生控制流逐样本判断，这类写法在eager模式下小批量测试可能不报错，但训练时开启静态图编译或者导出SavedModel格式时会直接失败，推理性能也会差一个数量级。
• 如果确实需要逐样本写复杂分支判断，优先用tf.vectorized_map做向量化映射，比逐样本循环效率高很多。

内容的提问来源于stack exchange，提问作者Philip Winchester