如何在Keras层间实现基于分支输出的自定义选择操作？

问题核心 #

Python原生的if-else是标量条件判断，无法处理Keras中的批量张量，也无法被TensorFlow的自动微分机制追踪，导致模型无法正常训练。必须使用TensorFlow提供的向量化条件操作来实现，比如tf.where，或者封装为自定义Keras层。

方案1：直接使用tf.where实现选择逻辑 #

tf.where(condition, x, y)会对张量中的每个元素进行判断：满足condition的位置取x对应的值，否则取y对应的值，完美适配批量张量的条件选择需求。

修改后的完整网络代码如下：

from tensorflow import keras
import tensorflow as tf
import numpy as np

# 定义网络输入
x1 = keras.Input(shape=(1), name="x1")
x2 = keras.Input(shape=(1), name="x2")

# 共享Dense层
shared_dense = keras.layers.Dense(20)
output_dense = keras.layers.Dense(1)

x11 = output_dense(shared_dense(x1))
x22 = output_dense(shared_dense(x2))

# 核心：用tf.where实现条件选择
# 条件：x11 >= x22，满足则选x1，否则选x2
Vm = tf.where(x11 >= x22, x1, x2)

# 后续输出处理
out = Vm - 0.5
out = keras.activations.sigmoid(out)

# 构建并编译模型
model = keras.Model([x1, x2], out)
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
    loss=tf.keras.losses.binary_crossentropy, 
    metrics=['accuracy']
)
model.summary()
tf.keras.utils.plot_model(model) # 可视化模型

方案2：封装为自定义Keras层（适合复杂逻辑） #

如果后续需要扩展选择逻辑，可以把条件选择封装成自定义Keras层，让代码更模块化：

class ConditionalSelector(keras.layers.Layer):
    def call(self, inputs):
        x1, x2, x11, x22 = inputs
        # 条件选择逻辑
        return tf.where(x11 >= x22, x1, x2)

# 在网络中使用自定义层
selector = ConditionalSelector()
Vm = selector([x1, x2, x11, x22])

训练数据集代码（无需修改） #

# 生成训练数据集
from scipy.stats import skewnorm

n=1000 # 每个类别样本数
s = 1 # 缩放输出范围
X1_0 = skewnorm.rvs(a = 0 ,loc=0, size=n)*s; X1_1 = skewnorm.rvs(a = 0 ,loc=1, size=n)*s
X2_0 = skewnorm.rvs(a = 0 ,loc=0, size=n)*s; X2_1 = skewnorm.rvs(a = 0 ,loc=1, size=n)*s

X1_train = list(X1_0) + list(X1_1)
X2_train = list(X2_0) + list(X2_1)

y_train = [x for x in (0,1) for i in range(0, n)]

# 调整为合适的格式
X1_train = np.array(X1_train).reshape(-1,1)
X2_train = np.array(X2_train).reshape(-1,1)
y_train = np.array(y_train)

# 训练模型
model.fit([X1_train, X2_train], y_train, epochs=10)

关键说明 #

• tf.where是向量化操作，会对批量中的每个样本独立判断，同时支持自动微分，完全兼容Keras的训练流程。
• 避免使用Python原生条件语句处理张量：这类语句是在图构建阶段执行，而非运行时对每个样本执行，无法实现动态选择逻辑，也会破坏反向传播链路。

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