Split output of a layer in keras

Question

Say, I have a layer with output dims (4, x, y). I want to split this into 4 separate (1, x, y) tensors, which I can use as input for 4 other layers.

What I'm essentially looking for is the opposite of the Merge layer. I know that there's no split layer in keras, but is there a simple way to do this in keras?

Answer 1

Are you looking for something like this?

import keras.backend as K
import numpy as np

val = np.random.random((4, 2, 3))
t = K.variable(value=val)
t1 = t[0, :, :]
t2 = t[1, :, :]
t3 = t[2, :, :]
t4 = t[3, :, :]

print('t1:\n', K.eval(t1))
print('t2:\n', K.eval(t2))
print('t3:\n', K.eval(t3))
print('t4:\n', K.eval(t4))
print('t:\n', K.eval(t))

It gives the following output:

t1:
 [[ 0.18787734  0.1085723   0.01127671]
 [ 0.06032621  0.14528386  0.21176969]]
t2:
 [[ 0.34292713  0.56848335  0.83797884]
 [ 0.11579451  0.21607392  0.80680907]]
t3:
 [[ 0.1908586   0.48186591  0.23439431]
 [ 0.93413448  0.535191    0.16410089]]
t4:
 [[ 0.54303145  0.78971165  0.9961108 ]
 [ 0.87826216  0.49061012  0.42450914]]
t:
 [[[ 0.18787734  0.1085723   0.01127671]
  [ 0.06032621  0.14528386  0.21176969]]

 [[ 0.34292713  0.56848335  0.83797884]
  [ 0.11579451  0.21607392  0.80680907]]

 [[ 0.1908586   0.48186591  0.23439431]
  [ 0.93413448  0.535191    0.16410089]]

 [[ 0.54303145  0.78971165  0.9961108 ]
  [ 0.87826216  0.49061012  0.42450914]]]

Note that, now t1, t2, t3, t4 is of shape(2,3).

print(t1.shape.eval()) # prints [2 3]

So, if you want to keep the 3d shape, you need to do the following:

t1 = t[0, :, :].reshape((1, 2, 3))
t2 = t[1, :, :].reshape((1, 2, 3))
t3 = t[2, :, :].reshape((1, 2, 3))
t4 = t[3, :, :].reshape((1, 2, 3))

Now, you get the spitted tensors in correct dimension.

print(t1.shape.eval()) # prints [1 2 3]

Hope that it will help you to solve your problem.

Answer 2

You can define Lambda layers to do the slicing for you:

from keras.layers import Lambda
from keras.backend import slice
.
.
x = Lambda( lambda x: slice(x, START, SIZE))(x)

For your specific example, try:

x1 = Lambda( lambda x: slice(x, (0, 0, 0), (1, -1, -1)))(x)
x2 = Lambda( lambda x: slice(x, (1, 0, 0), (1, -1, -1)))(x)
x3 = Lambda( lambda x: slice(x, (2, 0, 0), (1, -1, -1)))(x)
x4 = Lambda( lambda x: slice(x, (3, 0, 0), (1, -1, -1)))(x)