Выбор столбцов из трехмерного тензора в соответствии с одномерным тензором индексов (Tensorflow) - PullRequest
Новая
       7

Выбор столбцов из трехмерного тензора в соответствии с одномерным тензором индексов (Tensorflow)

0 голосов
/ 20 января 2019

Я ищу путь в тензорном потоке, учитывая два входа:

  1. input1, трехмерный тензор формы (batch_size, x, y)
  2. input2, a1D тензор формы (batch_size,), все значения которого находятся в диапазоне [0, y - 1] (включительно).

возвращает двумерный тензор формы (batch_size, x), такой что ith элемент в выходных данных равен столбцу input2[i]-th элемента ith в input1.

Пример: если input1 = [[[1,2], [3,4]], [[5,6], [7,8]], [[9,10], [11,12]]] (поэтому форма input1 равна (3, 2, 2)) и

input2 = [0, 1, 1], тогда я хочу получить вывод [[1,3], [6,8], [10,12]].

Объяснение: 0-й элемент в выводе равен [1,3], потому что 0-й элемент в input2 равен 0;таким образом, он становится 0-м столбцом в 0-м элементе input1.Последний элемент в выводе - [6,8], потому что последний элемент в input2 - 1;таким образом, он становится первым столбцом в последнем элементе input1.

Попытки:

Я попытался использовать tf.one_hot для достижения этой цели, (tf.reduce_sum(input1 * tf.one_hot(input2, y), 2)), но Tensorflow стал несчастным, когдаделая умножение, говоря: «ValueError: Размеры должны быть равны, но они равны 2 и 3 для« mul »(op:« Mul ») с входными формами: [3,2,2], [3,2].»

Любая помощь будет очень признательна, спасибо!

1 Ответ

0 голосов
/ 20 января 2019

Вы можете использовать tf.map_fn() для достижения этого. 

import tensorflow as tf
import numpy as np

input1 = [[[1,2], [3,4]], [[5,6], [7,8]], [[9,10], [11,12]]]
input2 = [0, 1, 1]

tf_input1 = tf.placeholder(shape=(None,2,2),dtype=tf.int32)
tf_input2 = tf.placeholder(shape=(None),dtype=tf.int32)

result = tf.map_fn(lambda x: x[0][:,x[1]], [tf_input1,tf_input2], dtype=tf.int32)

with tf.Session()as sess:
    result = sess.run(result,feed_dict={tf_input1:np.array(input1)
        ,tf_input2:np.array(input2)})
    print(result)

# print
[[ 1  3]
 [ 6  8]
 [10 12]]

Редактировать

tf.map_fn() медленный по сравнению с векторизованной операцией. Я добавил операцию умножения матриц. 

# shape= (3,2,1)
result = tf.cast(tf.expand_dims(tf.one_hot(input2, 2),-1),tf.int32)
# shape= (3,2)
result = tf.squeeze(tf.matmul(tf_input1, result))
...