Kerasで分岐を書いてみる

Kerasのバックエンドには便利な機能がいろいろありますが、その中にswitchというものがあります。
https://keras.io/ja/backend/#switch
これは以下のような使い方をして、分岐を書くものになっています。

switch(condition, then_expression, else_expression)

switchはfunction APIと組み合わせてモデル定義にも組み込むことができるようです。
ためしにやってみました。
https://github.com/shibuiwilliam/kerasswitch

import keras
from keras.models import Model
from keras.layers import Dense, Input, Lambda
from keras.optimizers import Adam
import keras.backend as K
from keras.utils.vis_utils import model_to_dot
from keras.utils import plot_model

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sn
import pydot
from IPython.display import SVG

# and data and or data
X = np.array([[0,0,0],[0,1,0],[1,0,0],[1,1,0],
              [0,0,1],[0,1,1],[1,0,1],[1,1,1]])
Y = np.array([0,0,0,1,0,1,1,1])
Y = keras.utils.to_categorical(Y, 2)

# input layer
inputs = Input(shape=X.shape[1:],name="input")

# and layers and or layers
def andGate(inputs):
    andDense = Dense(16, activation="relu", name="and1")(inputs)
    andDense = Dense(32, activation="relu", name="and2")(andDense) 
    return andDense
def orGate(inputs):
    orDense = Dense(32, activation="sigmoid", name="or1")(inputs)     
    return orDense

# switch cases for and and or
x_switch = Lambda(lambda x: K.switch(x[:,2]<1,
                                     andGate(x[:,:2]),
                                     orGate(x[:,:2])),
                  output_shape=(32,))(inputs)

# output layer
outputs = Dense(2, activation="softmax", name="softmax")(x_switch)

model = Model(inputs, outputs)

model.compile(loss='categorical_crossentropy',
              optimizer=Adam(),
              metrics=['accuracy'])
model.fit(X, Y,
          batch_size=4,
          epochs=20000,
          shuffle=True)

AndゲートとOrゲートを同居させたモデルで、switchでAndとOrを判定しています。
入力の形式は(8,3)です。
入力値の第3項は、Andの場合は0、Orの場合は1に設定し、switchはこの第3項を見て、第1,2項のみを次のレイヤの入力にします。
switchを使って分岐する部分はLambdaレイヤで囲みます。