Issue
I have two following tensors (note that they are both Tensorflow tensors which means they are still virtually symbolic at the time I construct the following slicing op before I launch a tf.Session()
):
params
: has shape (64,784, 256)indices
: has shape (64, 784)
and I want to construct an op that returns the following tensor:
output
: has shape (64,784) where
output[i,j] = params_tensor[i,j, indices[i,j] ]
What is the most efficient way in Tensorflow to do so?
ps: I tried with tf.gather
but couldn't make use of it to perform the operation I described above.
Many thanks.
-Bests
Solution
You can get exactly what you want using tf.gather_nd
. The final expression is:
tf.gather_nd(params, tf.stack([tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[0]), 1), [1, tf.shape(indices)[1]]), tf.transpose(tf.tile(tf.expand_dims(tf.range(tf.shape(indices)[1]), 1), [1, tf.shape(indices)[0]])), indices], 2))
This expression has the following explanation:
tf.gather_nd
does what you expected and uses the indices to gather the output from the paramstf.stack
combines three separate tensors, the last of which is the indices. The first two tensors specify the ordering of the first two dimensions (axis 0 and axis 1 of params/indices)- For the example provided, this ordering is simply 0, 1, 2, ..., 63 for axis 0, and 0, 1, 2, ... 783 for axis 1. These sequences are obtained with
tf.range(tf.shape(indices)[0])
andtf.range(tf.shape(indices)[1])
, respectively. For the example provided, indices has shape (64, 784). The other two tensors from the last point above need to have this same shape in order to be combined with
tf.stack
- First, an additional dimension/axis is added to each of the two sequences using
tf.expand_dims
. The use of
tf.tile
andtf.transpose
can be shown by example: Assume the first two axes of params and index have shape (5,3). We want the first tensor to be:[[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]]
We want the second tensor to be:
[[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]
These two tensors almost function like specifying the coordinates in a grid for the associated indices.
- First, an additional dimension/axis is added to each of the two sequences using
- For the example provided, this ordering is simply 0, 1, 2, ..., 63 for axis 0, and 0, 1, 2, ... 783 for axis 1. These sequences are obtained with
The final part of
tf.stack
combines the three tensors on a new third axis, so that the result has the same 3 axes as params.
Keep in mind if you have more or less axes than in the question, you need to modify the number of coordinate-specifying tensors in tf.stack
accordingly.
Answered By - Joe'
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.