Tensorflow tf.switch_case not working with keras Input

user_7182711 0
Viewed 194

I want to use tf.switch_case to be able to redirect learning flow for different branchs of the network using inputs however tf.switch_case does not work with Keras.Tensor ...

import tensorflow as tf
from tensorflow.keras.layers import Input
def f1(): return tf.constant(17)
def f2(): return tf.constant(31)
def f3(): return tf.constant(-1)
t_input = Input(shape=(1,), name="t_input")
r = tf.switch_case(t_input, branch_fns={0: f1, 1: f2}, default=f3)

Traceback (most recent call last):
  File "C:\Users\gen06917\PycharmProjects\BaysianTarnet\.venv\lib\site-packages\IPython\core\interactiveshell.py", line 3437, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-9-bd285228541c>", line 5, in <module>
    r = tf.switch_case(t_input, branch_fns={0: f1, 1: f2}, default=f3)
  File "C:\Users\gen06917\PycharmProjects\BaysianTarnet\.venv\lib\site-packages\tensorflow\python\ops\control_flow_ops.py", line 3616, in switch_case
    return _indexed_case_helper(branch_fns, default, branch_index, name)
  File "C:\Users\gen06917\PycharmProjects\BaysianTarnet\.venv\lib\site-packages\tensorflow\python\ops\control_flow_ops.py", line 3315, in _indexed_case_helper
    branch_fns, default, branch_index)
  File "C:\Users\gen06917\PycharmProjects\BaysianTarnet\.venv\lib\site-packages\tensorflow\python\ops\control_flow_ops.py", line 3249, in _indexed_case_verify_and_canonicalize_args
    type(branch_index)))
TypeError: branch_index must a Tensor, got <class 'tensorflow.python.keras.engine.keras_tensor.KerasTensor'>
keras tensorflow machine-learning python
1 Answers

I'm not sure why you'd do this in a Keras model. But just be careful when using these. For example, I'm not sure how well tested these are (for gradient flowing) if your f1 f2 are defining models in them. Nonetheless, you can do the following.

  1. Define your input as int32 type. Because tf.switch_case expects int32.
  2. Make sure you define batch_shape and not shape, so that you can index the input tensor to get the scalar value to be passed to tf.switch_case.
  3. Wrap the tf.switch_case in a Lambda layer.
import tensorflow as tf
from tensorflow.keras.layers import Input, Lambda

def f1(): return tf.constant(17)
def f2(): return tf.constant(31)
def f3(): return tf.constant(-1)

t_input = Input(batch_shape=(1,), dtype='int32', name="t_input")

r = Lambda(lambda x: tf.switch_case(x[0], branch_fns={0: f1, 1: f2}, default=f3))(t_input)
model = tf.keras.models.Model(inputs=t_input, outputs=r)

Using it

print(model.predict([1]))

Returns

31
user_16990750
Related