[FIXED] Randomly select layer in tensorflow model

Issue

I want to use different layers with specific probabilities in my network. Layers are the following classes.

class plus1(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, X):
        return X + 1    
    def compute_output_shape(self, batch_input_shape):
        return batch_input_shape

class plus2(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, X):
        return X + 2
    def compute_output_shape(self, batch_input_shape):
        return batch_input_shape

class plus3(keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    def call(self, X):
        return X + 3
    def compute_output_shape(self, batch_input_shape):
        return batch_input_shape

And the network is like below.

def f1():
    return plus1()
def f2():
    return plus2()
def f3():
    return plus3()

def simple_model(input_num):
    input_layer = Input(input_num)
    rand = tf.random.uniform((1,), minval=0, maxval=3, dtype=tf.int32)
    r = tf.switch_case(rand[0], branch_fns={0: f1, 1: f2, 2: f3})
    res = r(input_layer)
    model = Model(inputs=input_layer, outputs=res)
    return model

model = simple_model([1,])

Each time I run the code below, I get the same output, but I expected different ones. Is there any way to implement this?

model.predict([1])
>>> array([[4.]], dtype=float32)

Solution

It was the same problem I was facing and I found no solution. So I implemented distinct networks and then randomly selected from their output.

Answered By - Zahra Honjani