Keras conditional passing one model output to another model

Question

I am trying to neural network in keras, which first check if its a cat or dog (base model).

if it is a dog, then it is passes through another model (sub-model-1)

if it is a cat, then it is passes through another model (sub-model-2)

Sub-model are small model specially trained for classifying on the basis of breed.. so sub-model-1 will class dog into various dog's breed. . while sub-model-2 will classify cat into various cat's breed.

The problem i am facing is : I dont know how to add conditional layer,so that if base model has 5 million neuron and each sub-model has 2 million -2million neuron.. if an image is passed through base model, then it should only pass through sub-model1 or sub-model2.. so in total only 7 million neuron at action in passing one image to final output.

Any help, reference, everything would be appreciable.

Answer 1

Here is another solution which may train faster, run faster and use less RAM, give better performance, and be easier to use than the alternatives listed here.

Just use a single model with multiple outputs: a binary output (cat/dog), a cat breed output (multiclass), and a dog breed output (multiclass). During training, you can use a custom loss function to ignore the loss that corresponds to the wrong species (for example, ignore the cat breed output for dog images).

The benefits are:

• Faster training: just one model to train.
• Better performance: you can often get better performance when doing multi-task learning like this. That's because dog images and cat images have a lot in common, so it's helpful to train a single base neural network for both and then build specialized networks on top of that.
• Less RAM and less compute: instead of having to go through two CNNs (one for the cat/dog detector and one for the breed), we just go through one (the base network). This largely compensates the unnecessary computations that we do by going through the dog breed classifier even when the cat/dog detector says it's a cat.
• Easier to use: just one call to the model, and you get everything you need all at once. Moreover, if the cat/dog detector is not quite sure (for example it outputs a 50% probability), then you can at least have reasonable candidates for both cats and dogs.

Here's a working example. You just need to replace the data with your own. Note that there are three labels:

• cat/dog: for example [0, 1, 1, 0] for dog, cat, cat, dog
• cat breed: for example [-1, 2, 0, -1] for not-a-cat, 2nd cat class, 0th cat class, not-a-cat
• dog breed: for example [3, -1, -1, 1] for 3rd dog class, not-a-dog, not-a-dog, 1st dog class

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

np.random.seed(1)
tf.random.set_seed(1)

num_images = 200
num_cat_breeds = 10
num_dog_breeds = 15

X_train = np.random.random([num_images, 32, 32, 3])
y_breed = np.random.randint(num_cat_breeds + num_dog_breeds, size=num_images)
y_is_cat = y_breed < num_cat_breeds
y_cat_breed = np.where(y_is_cat, y_breed, -1)
y_dog_breed = np.where(y_is_cat, -1, y_breed - num_cat_breeds)

base_model = keras.Sequential([
    keras.layers.Conv2D(filters=32, kernel_size=3, activation="relu"),
    keras.layers.Flatten(),
])

model_is_cat = keras.Sequential([
    keras.layers.Dense(1, activation="sigmoid")
])

model_cat_breed = keras.Sequential([
    keras.layers.Dense(num_cat_breeds, activation="softmax")
])

model_dog_breed = keras.Sequential([
    keras.layers.Dense(num_dog_breeds, activation="softmax")
])

image_input = keras.layers.Input(shape=[32, 32, 3])
z = base_model(image_input)
is_cat = model_is_cat(z)
cat_breed = model_cat_breed(z)
dog_breed = model_dog_breed(z)
model = keras.Model(inputs=[image_input],
                    outputs=[is_cat, cat_breed, dog_breed])

def optional_crossentropy(y_true, y_pred):
    is_not_ignored = y_true != -1
    y_true_no_ignore = tf.where(is_not_ignored, y_true, 0)
    mask = tf.cast(is_not_ignored, tf.float32)
    return keras.losses.sparse_categorical_crossentropy(y_true_no_ignore, y_pred) * mask

model.compile(loss=["binary_crossentropy",
                    optional_crossentropy,
                    optional_crossentropy],
              optimizer="adam")
model.fit(X_train, [y_is_cat, y_cat_breed, y_dog_breed], epochs=2)

y_is_cat_pred, y_cat_breed_pred, y_dog_breed_pred = model.predict(X_train[:2])
print(y_is_cat_pred)
print(y_cat_breed_pred)
print(y_dog_breed_pred)