How to select rows from a 3-D Tensor in TensorFlow?

by

This is possible in TensorFlow, but slightly inconvenient, because tf.gather() currently only works with one-dimensional indices, and only selects slices from the 0th dimension of a tensor. However, it is still possible to solve your problem efficiently, by transforming the arguments so that they can be passed to tf.gather():

logits = ... # [2 x 4 x 4] tensor
indices = tf.constant([[0, 1], [1, 3]])

# Use tf.shape() to make this work with dynamic shapes.
batch_size = tf.shape(logits)[0]
rows_per_batch = tf.shape(logits)[1]
indices_per_batch = tf.shape(indices)[1]

# Offset to add to each row in indices. We use `tf.expand_dims()` to make 
# this broadcast appropriately.
offset = tf.expand_dims(tf.range(0, batch_size) * rows_per_batch, 1)

# Convert indices and logits into appropriate form for `tf.gather()`. 
flattened_indices = tf.reshape(indices + offset, [-1])
flattened_logits = tf.reshape(logits, tf.concat(0, [[-1], tf.shape(logits)[2:]]))

selected_rows = tf.gather(flattened_logits, flattened_indices)

result = tf.reshape(selected_rows,
                    tf.concat(0, [tf.pack([batch_size, indices_per_batch]),
                                  tf.shape(logits)[2:]]))

Note that, since this uses tf.reshape() and not tf.transpose(), it doesn’t need to modify the (potentially large) data in the logits tensor, so it should be fairly efficient.

More Related Contents:

  1. How to compile Tensorflow with SSE4.2 and AVX instructions?
  2. On Windows, running “import tensorflow” generates No module named “_pywrap_tensorflow” error
  3. Meaning of buffer_size in Dataset.map , Dataset.prefetch and Dataset.shuffle
  4. What is the meaning of the word logits in TensorFlow? [duplicate]
  5. How to set specific gpu in tensorflow?
  6. Can I use TensorBoard with Google Colab?
  7. What is the meaning of the “None” in model.summary of KERAS?
  8. Understanding TensorBoard (weight) histograms
  9. How do I split Tensorflow datasets?
  10. After building TensorFlow from source, seeing libcudart.so and libcudnn errors
  11. Specify connections in NN (in keras)
  12. What are the differences between all these cross-entropy losses in Keras and TensorFlow?
  13. Tensorflow: ImportError: libcusolver.so.8.0: cannot open shared object file: No such file or directory
  14. Tensorflow crashes with CUBLAS_STATUS_ALLOC_FAILED
  15. Keras Loss Function with Additional Dynamic Parameter
  16. TensorFlow: How to measure how much GPU memory each tensor takes?
  17. How to convert .pb to TFLite format?
  18. Create keras callback to save model predictions and targets for each batch during training
  19. Working with multiple graphs in TensorFlow
  20. How to extract data/labels back from TensorFlow dataset
  21. Obtaining output of an Intermediate layer in TensorFlow/Keras
  22. Making predictions with a TensorFlow model
  23. tf.data with multiple inputs / outputs in Keras
  24. Update only part of the word embedding matrix in Tensorflow
  25. In Keras, what exactly am I configuring when I create a stateful `LSTM` layer with N `units`?
  26. How to display custom images in TensorBoard using Keras?
  27. What is the difference between keras and tf.keras?
  28. feed data into a tf.contrib.data.Dataset like a queue
  29. parallelising tf.data.Dataset.from_generator
  30. how does tensorflow indexing work

Leave a Comment