Как использовать tf.gather в пакете?

Question

У меня есть тензор A = 10x1000 и тензор индекса B = 10x1000.Тензор B имеет значения в диапазоне 0-999 и используется для сбора значений из A (B[0,:] собирается из A[0,:], B[1,:] из A[1,:] и т. Д.).

Однако, еслиЯ использую tf.gather(A, B) Я получаю массив формы (10, 1000, 1000), когда я ожидаю 10x1000 тензор обратно.Любые идеи, как я мог бы это исправить?

РЕДАКТИРОВАТЬ

Допустим, A= [[1, 2, 3],[4,5,6]] и B = [[0, 1, 1],[2,1,0]] Что я хочу, чтобы иметь возможность выбрать A с использованием соответствующего BЭто должно привести к C = [[1, 2, 2],[6,5,4]].

Answer 1

1. Размеры тензоров известны заранее.

Сначала мы «расставляем» и параметры, и индексы (A и B соответственно) вдоль первого измерения. Затем мы применяем tf.gather() так, что строки A соответствуют строкам B. Наконец, мы складываем вместе результат.

import tensorflow as tf
import numpy as np

def custom_gather(a, b):
    unstacked_a = tf.unstack(a, axis=0)
    unstacked_b = tf.unstack(b, axis=0)
    gathered = [tf.gather(x, y) for x, y in zip(unstacked_a, unstacked_b)]
    return tf.stack(gathered, axis=0)

a = tf.convert_to_tensor(np.array([[1, 2, 3], [4, 5, 6]]), tf.float32)
b = tf.convert_to_tensor(np.array([[0, 1, 1], [2, 1, 0]]), dtype=tf.int32)

gathered = custom_gather(a, b)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(gathered))
# [[1. 2. 2.]
#  [6. 5. 4.]]

Для вас начальный случай с фигурами 1000x10 получаем:

a = tf.convert_to_tensor(np.random.normal(size=(10, 1000)), tf.float32)
b = tf.convert_to_tensor(np.random.randint(low=0, high=999, size=(10, 1000)), dtype=tf.int32)
gathered = custom_gather(a, b)
print(gathered.get_shape().as_list()) # [10, 1000]

Обновление

2. Первое измерение неизвестно (т. Е. None)

Предыдущее решение работает, только если первое измерение известно заранее. Если размер неизвестен, мы решаем его следующим образом:

• Мы складываем вместе два тензора, так что строки обоих тензоров складываются вместе:

# A = [[1, 2, 3], [4, 5, 6]]        [[[1 2 3]
#                            --->     [0 1 1]]
#                                    [[4 5 6]
# B = [[0, 1, 1], [2, 1, 0]]          [2 1 0]]]

• Мы перебираем элементы этого сложенного тензора (который состоит из сгруппированных вместе строк A и B) и, используя функцию tf.map_fn(), применяем tf.gather().

• Мы складываем обратно элементы, которые получаем с tf.stack()

import tensorflow as tf
import numpy as np

def custom_gather_v2(a, b):
    def apply_gather(x):
        return tf.gather(x[0], tf.cast(x[1], tf.int32))
    a = tf.cast(a, dtype=tf.float32)
    b = tf.cast(b, dtype=tf.float32)
    stacked = tf.stack([a, b], axis=1)
    gathered = tf.map_fn(apply_gather, stacked)
    return tf.stack(gathered, axis=0)

a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
b = np.array([[0, 1, 1], [2, 1, 0]], dtype=np.int32)

x = tf.placeholder(tf.float32, shape=(None, 3))
y = tf.placeholder(tf.int32, shape=(None, 3))

gathered = custom_gather_v2(x, y)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(gathered, feed_dict={x:a, y:b}))
# [[1. 2. 2.]
#  [6. 5. 4.]]