TensorFlow - индексирование тензоподобных тензоров

В numpy мы можем сделать это:

x = np.random.random((10,10))
a = np.random.randint(0,10,5)
b = np.random.randint(0,10,5)
x[a,b] # gives 5 entries from x, indexed according to the corresponding entries in a and b

Когда я попробую что-то эквивалентное в TensorFlow:

xt = tf.constant(x)
at = tf.constant(a)
bt = tf.constant(b)
xt[at,bt]

Последняя строка дает исключение "Плохой срез индекса индекса". Кажется, TensorFlow не поддерживает индексирование, например numpy или Theano.

Кто-нибудь знает, есть ли способ TensorFlow сделать это (индексирование тензора произвольными значениями). Я видел часть tf.nn.embedding, но я не уверен, что они могут быть использованы для этого, и даже если это возможно, это огромное обходное решение для чего-то такого простого.

(Сейчас я загружаю данные из x в качестве ввода и делаю индексирование в numpy, но я надеялся разместить x внутри TensorFlow, чтобы получить более высокую эффективность)

Ответ 1

Комментарий LDGN правильный. Это невозможно в данный момент и является запрошенной функцией. Если вы последуете за вопросом № 206 в github, вы будете обновлены, если/когда это будет доступно. Многие люди хотели бы эту функцию.

Ответ 2

Теперь вы можете сделать это с помощью tf.gather_nd. Скажем, у вас есть матрица m, как показано ниже:

| 1 2 3 4 |
| 5 6 7 8 |

И вы хотите построить матрицу r размера, допустим, 3x2, построенную из элементов m, например:

| 3 6 |
| 2 7 |
| 5 3 |
| 1 1 |

Каждый элемент r соответствует строке и столбцу m, и вы можете иметь матрицы rows и cols с этими индексами (основанные на нуле, поскольку мы программируем, а не делаем математику!)

       | 0 1 |         | 2 1 |
rows = | 0 1 |  cols = | 1 2 |
       | 1 0 |         | 0 2 |
       | 0 0 |         | 0 0 |

Что вы можете сделать в трехмерном тензоре:

| | 0 2 | | 1 1 | |
| | 0 1 | | 1 2 | |
| | 1 0 | | 2 0 | |
| | 0 0 | | 0 0 | |

Таким образом, вы можете получить от m до r через rows и cols следующим образом:

import numpy as np
import tensorflow as tf

m = np.array([[1, 2, 3, 4], [5, 6, 7, 8]])
rows = np.array([[0, 1], [0, 1], [1, 0], [0, 0]])
cols = np.array([[2, 1], [1, 2], [0, 2], [0, 0]])

x = tf.placeholder('float32', (None, None))
idx1 = tf.placeholder('int32', (None, None))
idx2 = tf.placeholder('int32', (None, None))
result = tf.gather_nd(x, tf.stack((idx1, idx2), -1))

with tf.Session() as sess:
    r = sess.run(result, feed_dict={
        x: m,
        idx1: rows,
        idx2: cols,
    })
print(r)

Вывод:

[[ 3.  6.]
 [ 2.  7.]
 [ 5.  3.]
 [ 1.  1.]]

Ответ 3

При tf 0.11 было выполнено базовое индексирование. Более продвинутое индексирование (например, булевское индексирование) по-прежнему отсутствует, но, по-видимому, планируется для версий futur.

Расширенное индексирование можно отслеживать с помощью https://github.com/tensorflow/tensorflow/issues/4638