Can I use loops inside a model using functional API?

Question

I have a trained keras model which takes inputs of size (batchSize,2). This works well and gives good results.

My main problem is to have a model which takes an input a vector of size(batchSize,2,16) and slice it inside the model to 16 vectors of size(batchSize,2) and concatenate the outputs together.

I have used this code for this

    y = layers.Input(shape=(2,16,))

    model_x= load_model('saved_model')

    for i in range(16):
        x_input = Lambda(lambda x: x[:, :, i])(y)

        if i == 0:
           x_output = model_x(x_input)
        else:
            x_output = layers.concatenate([x_output, 
                                      model_x(x_input)])

    x_output = Lambda(lambda x: x[:, :tf.cast(N, tf.int32)])(x_output)

    final_model = Model(y, x_output)

Although the saved model gives me good performance, this code does not trains well and doesn't give the intended performance. What can I do to get better results?

Answer 1

I can't say anything about the bad performance of your final model because it might be due to various reasons and this is not readily evident from the content of your question. But to answer your original question: yes, you can use for loops that way, because you are essentially creating layers/tensors and connecting them to each other (i.e. building the graph of the model). So it's a valid thing to do. The problem might be somewhere else, e.g. a wrong indexing, a wrong loss function, etc.

Further, you can build your final model in a much simpler approach. You already have a trained model which gets inputs of shape (batch_size, 2) and gives outputs of shape (batch_size, 8). Now you want to build a model which takes inputs of shape (batch_size, 2, 16), apply the already trained model on each of the 16 (batch_size, 2) segments and then concatenate the results. You can easily do that with a TimeDistributed wrapper:

# load your already trained model
model_x = load_model('saved_model')

inp = layers.Input(shape=(2,16))
# this makes the input shape as `(16,2)`
x = layers.Permute((2,1))(inp)
# this would apply `model_x` on each of the 16 segments; the output shape would be (None, 16, 8)
x = layers.TimeDistributed(model_x)(x)
# flatten to make it have a shape of (None, 128)
out = layers.Flatten()(x)

final_model = Model(inp, out)