Keras: Is there any workaround to split the output of an intermediate layer without using Lamda layer?

Question

Say, I have a 10x10x4 intermediate output of a convolution layer, which I need to split into 100 1x1x4 volume and apply softmax on each to get 100 outputs from the network. Is there any way to accomplish this without using the Lambda layer? The issue with the Lambda layer in this case is this simple task of splitting takes 100 passes through the lambda layer during forward pass, which makes the network performance very slow for my practical use. Please suggest a quicker way of doing this.

Edit: I had already tried the Softmax+Reshape approach before asking the question. With that approach, I would be getting a 10x10x4 matrix reshaped to a 100x4 Tensor with use of Reshape as the output. What I really need is a multi output network with 100 different outputs. In my application, it is not possible to jointly optimize over the 10x10 matrix, but I get good results by using a network with 100 different outputs with the Lambda layer.

Here are code snippets of my approach using the Keras functional API:

With Lambda layer (slow, gives 100 Tensors of shape (None, 4) as desired):

# Assume conv_output is output from a convolutional layer with shape (None, 10, 10,4)
preds = []
for i in range(10):
    for j in range(10):
        y = Lambda(lambda x, i,j: x[:, i, j,:],  arguments={'i': i,'j':j})(conv_output)
        preds.append(Activation('softmax',name='predictions_' + str(i*10+j))(y))

model = Model(inputs=img, outputs=preds, name='model')
model.compile(loss='categorical_crossentropy',
              optimizer=Adam(),
              metrics=['accuracy']

With Softmax+Reshape (fast, but gives Tensor of shape (None, 100, 4))

# Assume conv_output is output from a convolutional layer with shape (None, 10, 10,4)
y = Softmax(name='softmax', axis=-1)(conv_output)
preds = Reshape([100, 4])(y)
model = Model(inputs=img, outputs=preds, name='model')
model.compile(loss='categorical_crossentropy',
              optimizer=Adam(),
              metrics=['accuracy']

I don't think in the second case it is possible to individually optimize over each of the 100 outputs (probably one can think of it as learning the joint distribution, whereas I need to learn the marginals as in the first case). Please let me know if there is any way to accomplish what I am doing with the Lambda layer in the first code snippet in a faster way

Answer 1

You can use the Softmax layer and set the axis argument to the last axis (i.e. -1) to apply softmax over that axis:

from keras.layers import Softmax

soft_out = Softmax(axis=-1)(conv_out)

Note that the axis argument by default is set to -1, so you may not even need to pass that.