Python – Unable to merge keras models after adding a Lambda layer

I have a list of 3 keras models, each with an output shape of (None, 2). I also have a generic keras base model that generates inputs to them. My goal is to combine these 4 models, but only get the first output from each model in the list (so the shape of the final output should be (None, 3)). I had a problem when I tried using Lambda layers to extract the first output from each model.

If I omit the Lambda step and simply combine the models as follows, it creates a model that provides the correct output of shape (None, 6):

>>> sequentials = [Sequential([base_model, m]) for m in models]
>>> output = merge([s.output for s in sequentials], mode='concat')
>>> combined = Model(input=base_model.layers[0].input, output=output)
>>> combined.predict(X)
array([[  8.52127552e-01,   1.47872433e-01,   1.89960217e-13,
          1.00000000e+00,   7.56258190e-01,   2.43741751e-01]], dtype=float32)

The problem arises when I first use a Lambda layer to extract the first value from each model:

>>> print([m.output_shape for m in models])
[(None, 2), (None, 2), (None, 2)]
>>> for m in models:
        m.add(Lambda(lambda x: x[0], output_shape=(1,)))
>>> print([m.output_shape for m in models])
[(None, 1), (None, 1), (None, 1)]
>>> sequentials = [Sequential([base_model, m]) for m in models]
>>> print([s.output_shape for s in sequentials])
[(None, 1), (None, 1), (None, 1)]
>>> output = merge([s.output for s in sequentials],
                   output_shape=(len(sequentials),), mode='concat')
>>> combined = Model(base_model.layers[0].input, output=output)
>>> print(combined.output_shape)
(None, 3)
>>> combined.predict(X)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-3-4f4ed3bd605d> in <module>()
----> 1 ann.combined.predict(X)

./.virtualenvs/py3/lib/python3.4/site-packages/keras/engine/training.py in predict(self, x, batch_size, verbose)
   1217         f = self.predict_function
   1218         return self._predict_loop(f, ins,
-> 1219                                   batch_size=batch_size, verbose=verbose)
   1220
   1221     def train_on_batch(self, x, y,

./.virtualenvs/py3/lib/python3.4/site-packages/keras/engine/training.py in _predict_loop(self, f, ins, batch_size, verbose)
    904
    905             for i, batch_out in enumerate(batch_outs):
--> 906                 outs[i][batch_start:batch_end] = batch_out
    907             if verbose == 1:
    908                 progbar.update(batch_end)

ValueError: could not broadcast input array from shape (6) into shape (1)

What is the correct way to combine each model while only getting a single output value from them?

Note that if I apply the Lambda function after merging the models, I can use it successfully, as follows:

>>> sequentials = [Sequential([base_model, m]) for m in models]
>>> output = merge([s.output for s in sequentials], mode='concat')
>>> filtered = Lambda(lambda x: x[:,::2], lambda s: (s[-1] / 2,))(output)
>>> combined = Model(input=base_model.layers[0].input, output=filtered)
>>> combined.predict(X)
array([[  1.89960217e-13,   7.56258249e-01,   8.52127552e-01]], type=float32)

But I would like to know how to apply it before merging.

The problem is a little inconsistency in the Lambda slice. Although the shape of the output does not take into account the batch dimension – we should remember that the tensor provides this extra dimension to the Lambda layer as well. This is why the following line causes the error:

m.add(Lambda(lambda x: x[0], output_shape=(1,)))

This should be changed to:

m.add(Lambda(lambda x: x[:,:1], output_shape=(1,))) 

Note the following slicing method:

m.add(Lambda(lambda x: x[:,0], output_shape=(1,)))

Because it changes the dimension of the tensor.