Keras: gather tensor changes batch dimension

user_886104 0
Viewed 1268

I have an input tensor of shape (5, 2), representing five points in 2D space.

I would like to take the first point, then subtract it from all five points.

Reading around, I figured I can use K.gather to slice and repeat the first layer.

After applying it in a Lambda layer the batch dimension is overwritten:

_input = Input(shape=(5, 2))
x = Reshape((5 * 2,))(_input)
x_ = Lambda(lambda t: K.gather(t, [0, 1] * 5))(x)

results in:

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            (None, 5, 2)         0                                            
__________________________________________________________________________________________________
reshape_1 (Reshape)             (None, 10)           0           input_1[0][0]                    
__________________________________________________________________________________________________
lambda_1 (Lambda)               (10, 10)             0           reshape_1[0][0]                  
__________________________________________________________________________________________________

What am I doing wrong?

Also, is there a simpler way to do this?

keras tensorflow python
2 Answers

gather function is returning values of provided indices from batch(0th)-axis. Therefore, it's giving us a list (length=10) of 1st (index:0) and 2nd (index:1) samples (shape (10,)) from the batch with the shape of (10, 10) whereas we want 1st (index:0) and 2nd (index:1) feature point of each samples in the batch. To handle this, we can transpose the tensor before using gather function so that gather function picks right values and finally resulted tensor should be transposed again.

_input = Input(shape=(5, 2))
x = Reshape((5 * 2,))(_input)
x_ = Lambda(lambda t: K.transpose(K.gather(K.transpose(t), [0, 1]*5)))(x)

Output:

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 5, 2)]            0         
_________________________________________________________________
reshape (Reshape)            (None, 10)                0         
_________________________________________________________________
lambda (Lambda)              (None, 10)                0         
=================================================================
user_51365500

If you use tf.gather(), you can avoid to use transpose operation as described by @bit01. There is an axis argument in tf.gather().

_input = Input(shape=(5, 2))
x = Reshape((5 * 2,))(_input)
x_ = Lambda(lambda t: tf.gather(t, [0, 1]*5, axis=1))(x)

# Layer (type)                 Output Shape              Param #   
# =================================================================
# input_2 (InputLayer)         (None, 5, 2)              0         
# _________________________________________________________________
# reshape_2 (Reshape)          (None, 10)                0         
# _________________________________________________________________
# lambda_1 (Lambda)            (None, 10)                0         
# =================================================================
user_72123650
Related