Keras条件分支模型实现求助：基础模型输出定向至对应子模型

Hey there, I’ve helped folks with similar conditional routing setups in Keras before—let’s break this down so you can get your model working exactly how you want it.

The core challenge here is making sure only one sub-model runs after the base classifier, instead of both. That way, you only activate the base model’s 5M neurons plus the relevant sub-model’s 2M, hitting your ~7M total per image target. Let’s dive into the implementation:

Step 1: Define Your Base and Sub-Models #

First, we’ll build the base cat/dog classifier (which also extracts reusable features) and the two breed-classification sub-models. Make sure the sub-models accept the same feature shape that the base model outputs—this keeps the data flow clean.

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

# Base model: Cat/dog classifier + feature extractor
def build_base_model(input_shape):
    inputs = layers.Input(shape=input_shape)
    # Replace this with your actual 5M-neuron architecture (e.g., fine-tuned ResNet)
    x = layers.Conv2D(32, (3,3), activation='relu')(inputs)
    x = layers.MaxPooling2D()(x)
    x = layers.Conv2D(64, (3,3), activation='relu')(x)
    x = layers.MaxPooling2D()(x)
    x = layers.Conv2D(128, (3,3), activation='relu')(x)
    x = layers.MaxPooling2D()(x)
    features = layers.GlobalAveragePooling2D()(x)  # Shared feature vector
    cat_dog_output = layers.Dense(1, activation='sigmoid', name='cat_dog_cls')(features)
    return Model(inputs=inputs, outputs=[features, cat_dog_output], name='base_model')

# Sub-model 1: Dog breed classifier (2M neurons)
def build_dog_breed_model(feature_dim):
    inputs = layers.Input(shape=(feature_dim,))
    x = layers.Dense(256, activation='relu')(inputs)
    x = layers.Dropout(0.5)(x)
    outputs = layers.Dense(15, activation='softmax', name='dog_breed')(x)  # Adjust num breeds
    return Model(inputs=inputs, outputs=outputs, name='dog_breed_model')

# Sub-model 2: Cat breed classifier (2M neurons)
def build_cat_breed_model(feature_dim):
    inputs = layers.Input(shape=(feature_dim,))
    x = layers.Dense(256, activation='relu')(inputs)
    x = layers.Dropout(0.5)(x)
    outputs = layers.Dense(12, activation='softmax', name='cat_breed')(x)  # Adjust num breeds
    return Model(inputs=inputs, outputs=outputs, name='cat_breed_model')

Step 2: Add Conditional Routing with tf.cond() #

The key here is using TensorFlow’s tf.cond() to route the base model’s features to the correct sub-model. This ensures only one sub-model is executed per image—no wasted neuron activations. We’ll wrap this logic in a Lambda layer to integrate it with Keras’ functional API.

# Initialize models
input_shape = (224, 224, 3)  # Adjust to your image size
base_model = build_base_model(input_shape)
feature_dim = base_model.output[0].shape[-1]  # Get size of feature vector

dog_model = build_dog_breed_model(feature_dim)
cat_model = build_cat_breed_model(feature_dim)

# Define routing logic
def route_to_submodel(features, cat_dog_pred):
    # Check if prediction is dog (sigmoid > 0.5)
    is_dog = tf.greater(cat_dog_pred, 0.5)
    
    # Define functions for each route
    def run_dog_model():
        return dog_model(features)
    
    def run_cat_model():
        return cat_model(features)
    
    # Use tf.cond to execute only the relevant branch
    breed_pred = tf.cond(is_dog, run_dog_model, run_cat_model)
    return breed_pred

# Assemble the full model
inputs = layers.Input(shape=input_shape)
features, cat_dog_pred = base_model(inputs)
breed_pred = layers.Lambda(
    lambda x: route_to_submodel(x[0], x[1]),
    name='breed_router'
)([features, cat_dog_pred])

full_model = Model(inputs=inputs, outputs=[cat_dog_pred, breed_pred], name='multi_stage_model')

Step 3: Train the Model with Multi-Output Loss #

When training, you’ll need two sets of labels: one for cat/dog classification, and one for breed classification (only the relevant breed label matters per image). The tf.cond() logic ensures that only the active sub-model’s loss contributes to training.

full_model.compile(
    optimizer='adam',
    loss={
        'cat_dog_cls': 'binary_crossentropy',
        'breed_router': 'sparse_categorical_crossentropy'  # Use this if labels are integers
    },
    metrics={
        'cat_dog_cls': 'accuracy',
        'breed_router': 'accuracy'
    }
)

# Example training call (replace with your dataset)
# full_model.fit(
#     x_train,
#     {'cat_dog_cls': y_cat_dog_labels, 'breed_router': y_breed_labels},
#     epochs=15,
#     batch_size=32,
#     validation_split=0.2
# )

Key Notes to Ensure Correct Activation #

• No wasted computation: tf.cond() doesn’t execute the unused branch in static graph mode (the default in TensorFlow/Keras), so only the relevant sub-model’s neurons are activated.
• Feature reuse: By having the base model output both features and the cat/dog prediction, you avoid reprocessing the image twice.
• Flexibility: If you need more complex routing (e.g., threshold adjustments), you can modify the is_dog condition in the routing function.

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