tensorflow: slicing a tensor along the second dimension

0
votes

I have a tensor X whose shape is (None, 56, 300, 1), and another tensor y whose shape is (None, 15), the first dimension of these tensors is batch_size, I wanna use y as index to get a tensor z, the shape of z is (None, 15, 300, 1). Is there any decent way to do this?

I write a simple code to test, for I found it's difficult for me because in practice I don't know the batch_size(first dimension of these tensors is None),

Here is my test code:

import numpy as np
import tensorflow as tf

# In this test code , batch_size is 4.
# params' shape is (4, 3, 2 ,1), in practice is (None, 56, 300, 1), 
params = [
            [[['a0'], ['b0']], [['d0'], ['e0']], [['f0'], ['g0']]], 
            [[['a1'], ['b1']], [['d1'], ['e1']], [['f1'], ['g1']]],
            [[['a2'], ['b2']], [['d2'], ['e2']], [['f2'], ['g2']]],
            [[['a3'], ['b3']], [['d3'], ['e3']], [['f3'], ['g3']]],
        ]

# ind's shape is (4, 2) (In practice is (None, 15)), 
# so I wanna get output whose's shape is (4, 2, 2, 1), (In practice is (None, 15, 300, 1))
ind = [[1, 0], [0, 2], [2, 0], [2, 1]]
#ouput = [
# [[['d0'], ['e0']], [['a0'], ['b0']]],
# [[['a1'], ['b1']], [['f1'], ['g1']]],
# [[['f2'], ['g2']], [['a2'], ['b2']]],
# [[['f3'], ['g3']], [['d3'], ['e3']]]
#]

with tf.variable_scope('gather') as scope:
    tf_par = tf.constant(params)
    tf_ind = tf.constant(ind)
    res = tf.gather_nd(tf_par, tf_ind)

with tf.Session() as sess:
    init = tf.global_variables_initializer()
    print sess.run(res)
    print res
2
pythontensorflow
could you add your expectation result of the sample code – Yuwen Yan
confused here, to my understanding [1, 0] should be [['a1'], ['b1']], but your expectation is [['d0'], ['e0']] – Yuwen Yan
@YuwenYan [[['a0'], ['b0']], [['d0'], ['e0']], [['f0'], ['g0']]] is the first sample. [1, 0] responses to the first sample, so the answer should be [['d0'], ['e0']], [['a0'], ['b0']]. Note that length of ind is 4, each element response to each sample in params – Joey

2 Answers

3
votes

To slice x along the second dimension with ind, that is, to slice

  • tensor x of shape (d0, d1, d2,...), d0 being possibly None,
  • with a tensor of indices ind of shape (d0, n1),
  • to obtain a tensor y of shape (d0, n1, d2, ...),

you could use tf.gather_nd along with tf.shape to get the shape at run time:

ind_shape = tf.shape(ind)
ndind = tf.stack([tf.tile(tf.range(ind_shape[0])[:, None], [1, ind_shape[1]]),
                  ind], axis=-1)
y = tf.gather_nd(x, ndind)
0
votes

For results you suppose, you should use: 

ind = [[0, 1], [0, 0], [1, 0], [1, 2], [2, 2], [2, 0], [3, 2], [3, 1]]

Update

You can use this code for get what you want, with current input:

with tf.variable_scope('gather') as scope:
    tf_par = tf.constant(params)
    tf_ind = tf.constant(ind)

    tf_par_shape = tf.shape(tf_par)
    tf_ind_shape = tf.shape(tf_ind)
    tf_r = tf.div(tf.range(0, tf_ind_shape[0] * tf_ind_shape[1]), tf_ind_shape[1])
    tf_r = tf.expand_dims(tf_r, 1)
    tf_ind = tf.expand_dims(tf.reshape(tf_ind, shape = [-1]), 1)
    tf_ind = tf.concat([tf_r, tf_ind], axis=1)

    res = tf.gather_nd(tf_par, tf_ind)
    res = tf.reshape(res, shape = (-1, tf_ind_shape[1], tf_par_shape[2], tf_par_shape[3]))