<h2>The problem with conditionals in neural networks</h2>
<p>The issue with a switch or conditionals (like if-then-else) as part of a neural network is that conditionals are not differentiable everywhere. Therefore the automatic differentiation methods would not work directly and solving this is super complex. Check <a href="https://cs.stackexchange.com/a/70620">this</a> for more details.</p>
<p>A shortcut is you can end up training 3 separate models independently, and then during inference uses a control flow of conditionals to infer from them.</p>
<pre><code>#Training - 
model1 = model.fit(all images, P(cat/dog))
model2 = model.fit(all images, P(cat))
model3 = model.fit(all images, P(dog))
final prediction = argmax(model2, model3)

#Inference - 
if model1.predict == Cat: 
    model2.predict
else:
    model3.predict
</code></pre>
<p>But I don't think you are looking for that. <strong>I think you are looking to include conditionals as part of the computation graph itself.</strong></p>
<p>Sadly, there is no direct way for you to build an if-then condition as part of a computation graph as per my knowledge. The <code>keras.switch</code> that you see allows you to work with tensor outputs but not with layers of a graph during training. That's why you will see it being used as part of loss functions and not in computation graphs (throws input errors).</p>
<h2>A possible Solution - Skip connections &amp; soft-switching</h2>
<p>You can, however, try to build something similar with <code>skip connections</code> and <code>soft switching</code>.</p>
<p>A skip connection is a connection from a previous layer to another layer that allows you to pass information to the subsequent layers. This is quite common in very deep networks where information from the original data is subsequently lost. Check <a href="https://arxiv.org/pdf/1505.04597.pdf" rel="noreferrer">U-net</a> or <a href="https://arxiv.org/pdf/1512.03385.pdf" rel="noreferrer">Resnet</a> for example, which uses skip connections between layers to pass information to future layers.</p>
<p><a href="https://i.stack.imgur.com/9sroF.png" rel="noreferrer"><img src="https://i.stack.imgur.com/9sroF.png" alt="enter image description here" /></a></p>
<p>The next issue is the issue of switching. You want to switch between 2 possible paths in the graph. What you can do is a soft-switching method which I took as inspiration from <a href="https://arxiv.org/pdf/1905.08743.pdf" rel="noreferrer">this paper</a>. Notice that in order to <code>switch</code> between 2 distribution of words (one from the decoder and another from the input), the authors multiply them with <code>p</code> and <code>(1-p)</code> to get a cumulative distribution. This is a soft-switch that allows the model to pick the next predicted word from either the decoder or from the input itself. (helps when you want your chatbot to speak the words that were input by the user as part of its response to them!)</p>
<p><a href="https://i.stack.imgur.com/Gpv0C.png" rel="noreferrer"><img src="https://i.stack.imgur.com/Gpv0C.png" alt="enter image description here" /></a></p>
<p>With an understanding of these 2 concepts, let's try to intuitively build our architecture.</p>
<ol>
<li><p>First we need a single-input multi-output graph since we are training 2 models</p>
</li>
<li><p>Our first model is a multi-class classification that predicts individual probabilities for Cat and Dog separately. This will be trained with the activation of <code>softmax</code> and a <code>categorical_crossentropy</code> loss.</p>
</li>
<li><p>Next, let's take the logit which predicts the probability of Cat, and multiply the convolution layer 3 with it. This can be done with a <code>Lambda</code> layer.</p>
</li>
<li><p>And similarly, let's take the probability of Dog and multiply it with the convolution layer 2. This can be seen as the following -</p>
<ul>
<li>If my first model predicts a cat and not a dog, perfectly, then the computation will be <code>1*(Conv3)</code> and <code>0*(Conv2)</code>.</li>
<li>If the first model predicts a dog and not a cat, perfectly, then the computation will be <code>0*(Conv3)</code> and <code>1*(Conv2)</code></li>
<li>You can think of this as either a <code>soft-switch</code> OR a <code>forget gate</code> from <a href="https://colah.github.io/posts/2015-08-Understanding-LSTMs/" rel="noreferrer">LSTM</a>. The <code>forget gate</code> is a sigmoid (0 to 1) output that multiplies the cell state to gate it and allow the LSTM to forget or remember previous time-steps. Similar concept here!</li>
</ul>
</li>
<li><p>These Conv3 and Conv2 can now be further be processed, flattened, concatenated, and passed to another Dense layer for the final prediction.</p>
</li>
</ol>
<p>This way if the model is not sure about a dog or a cat, both conv2 and conv3 features participate in the second model's predictions. This is how you can use <code>skip connections</code> and <code>soft switch</code> inspired mechanism to add some amount of conditional control flow to your network.</p>
<p>Check my implementation of the computation graph below.</p>
<pre><code>from tensorflow.keras import layers, Model, utils
import numpy as np

X = np.random.random((10,500,500,3))
y = np.random.random((10,2))

#Model
inp = layers.Input((500,500,3))

x = layers.Conv2D(6, 3, name='conv1')(inp)
x = layers.MaxPooling2D(3)(x)

c2 = layers.Conv2D(9, 3, name='conv2')(x)
c2 = layers.MaxPooling2D(3)(c2)

c3 = layers.Conv2D(12, 3, name='conv3')(c2)
c3 = layers.MaxPooling2D(3)(c3)

x = layers.Conv2D(15, 3, name='conv4')(c3)
x = layers.MaxPooling2D(3)(x)

x = layers.Flatten()(x)
out1 = layers.Dense(2, activation='softmax', name='first')(x)

c = layers.Lambda(lambda x: x[:,:1])(out1)
d = layers.Lambda(lambda x: x[:,1:])(out1)

c = layers.Multiply()([c3, c])
d = layers.Multiply()([c2, d])

c = layers.Conv2D(15, 3, name='conv5')(c)
c = layers.MaxPooling2D(3)(c)
c = layers.Flatten()(c)

d = layers.Conv2D(12, 3, name='conv6')(d)
d = layers.MaxPooling2D(3)(d)
d = layers.Conv2D(15, 3, name='conv7')(d)
d = layers.MaxPooling2D(3)(d)
d = layers.Flatten()(d)

x = layers.concatenate([c,d])
x = layers.Dense(32)(x)
out2 = layers.Dense(2, activation='softmax',name='second')(x)

model = Model(inp, [out1, out2])
model.compile(optimizer='adam', loss='categorical_crossentropy', loss_weights=[0.5, 0.5])

model.fit(X, [y, y], epochs=5)

utils.plot_model(model, show_layer_names=False, show_shapes=True)
</code></pre>
<pre><code>Epoch 1/5
1/1 [==============================] - 1s 1s/step - loss: 0.6819 - first_loss: 0.7424 - second_loss: 0.6214
Epoch 2/5
1/1 [==============================] - 0s 423ms/step - loss: 0.6381 - first_loss: 0.6361 - second_loss: 0.6400
Epoch 3/5
1/1 [==============================] - 0s 442ms/step - loss: 0.6137 - first_loss: 0.6126 - second_loss: 0.6147
Epoch 4/5
1/1 [==============================] - 0s 434ms/step - loss: 0.6214 - first_loss: 0.6159 - second_loss: 0.6268
Epoch 5/5
1/1 [==============================] - 0s 427ms/step - loss: 0.6248 - first_loss: 0.6184 - second_loss: 0.6311
</code></pre>
<p><a href="https://i.stack.imgur.com/8JTb2.png" rel="noreferrer"><img src="https://i.stack.imgur.com/8JTb2.png" alt="enter image description here" /></a></p>
CNN model conditional layer in Keras
<p>I am trying to build a <code>conditional CNN</code> model. The model is,</p>
<p><a href="https://i.stack.imgur.com/ww0dP.png" rel="noreferrer"><img src="https://i.stack.imgur.com/ww0dP.png" alt="enter image description here" /></a></p>
<p>At the <code>first stage</code> of my model, I feed my data to <code>Model 1</code> then, <code>based on the prediction of Model 1</code>, I want to <code>train the model to Conditional Cat model or Conditional Dog model</code> and finally, give the output from Conditional Cat model or Conditional Dog model. <strong>How Can I do this?</strong></p>
<p><strong>Note:</strong>
My effort is,</p>
<pre><code>import keras
from keras.layers import *
from keras.models import *
from keras.utils import *

img_rows,img_cols,number_of_class = 256,256,2
input = Input(shape=(img_rows,img_cols,3))

#----------- main model (Model 1) ------------------------------------
conv_01 = Convolution2D(64, 3, 3, activation='relu',name = 'conv_01') (input)
conv_02 = Convolution2D(64, 3, 3, activation='relu',name = 'conv_02') (conv_01)

skip_dog =  conv_02

conv_03 = Convolution2D(64, 3, 3, activation='relu',name = 'conv_03') (conv_02)

skip_cat =  conv_03

conv_04 = Convolution2D(64, 3, 3, activation='relu',name = 'conv_04') (conv_03)


flatten_main_model =  Flatten() (conv_04)
Output_main_model = Dense(units = number_of_class , activation = 'softmax', name = &quot;Output_layer&quot;)(flatten_main_model)


#----------- Conditional  Cat model ------------------------------------ 
conv_05 = Convolution2D(64, 3, 3, activation='relu',name = 'conv_05') (skip_cat)
flatten_cat_model =  Flatten() (conv_05)
Output_cat_model = Dense(units = number_of_class , activation = 'softmax', name = &quot;Output_layer_cat&quot;)(flatten_cat_model)

#----------- Conditional  Dog model ------------------------------------ 
conv_06 = Convolution2D(64, 3, 3, activation='relu',name = 'conv_06') (skip_dog)
flatten_dog_model =  Flatten() (conv_06)
Output_dog_model = Dense(units = number_of_class , activation = 'softmax', name = &quot;Output_layer_dog&quot;)(flatten_dog_model)

#----------------------------- My discrete 3 models --------------------------------
model_01 = Model(inputs = input , outputs = Output_main_model,name = 'model_main')
model_02_1 = Model(inputs = input , outputs = Output_cat_model ,name = 'Conditional_cat_model')
model_02_2 = Model(inputs = input , outputs = Output_dog_model ,name = 'Conditional_dog_model')
</code></pre>
<p>How can I merge these 3 models (<code>model_01, model_02_1, model_02_2</code>) based on these conditions?</p>
<p>**Conditions are: **</p>
<ol>
<li>Feed data to model <code>model_01</code></li>
<li>Based on <code>model_01</code> result feed data to <code>model_02_1 or model_02_2</code></li>
<li>Next, predict the final output from <code>model_02_1 or model_02_2</code></li>
</ol>