Блог

У меня есть 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]] .

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.]]
  

Использовать tf.gather с batch_dims=-1 :

 import numpy as np
import tensorflow as tf

rois = np.array([[1, 2, 3],[3, 2, 1]])
ind = np.array([[0, 2, 1, 1, 2, 0, 0, 1, 1, 2],
        [0, 1, 2, 0, 2, 0, 1, 2, 2, 2]])
tf.gather(rois, ind, batch_dims=-1)
# output:
# <tf.Tensor: shape=(2, 10), dtype=int64, numpy=
# array([[1, 3, 2, 2, 3, 1, 1, 2, 2, 3],
#       [3, 2, 1, 3, 1, 3, 2, 1, 1, 1]])>