根据两个索引向量以及Tensorflow中的一维从张量中提取/收集元素

Question

很抱歉，问题标题冗长乏味，难以理解。基本上，我想在 tensorflow 中实现一个函数：

例如 对于维度为 [10, 10, 7, 1] 的张量 A 和索引矩阵 B = array([[1,3,5],[2,4,6]])。我想根据B每行的索引提取A中的元素以及axis = 2（遵循Python约定，A有0,1,2,3四个轴）。

所以示例的结果应该是一个维度为 [10, 10, 3, 2] 的张量 C，其中第三个维度是由于根据索引 [1,3,5] 沿轴 = 2 选择 A 中的元素] 或 [2,4,6]，并且第四维等于 B 的第一个维度（即此处 B 的行数），因为我们在这里沿该维度进行了两次选择。

在张量流中实现这一点的任何“张量青睐”线索，而不是分两步完成？我没有看到使用 tf.gather_nd() 或 tf.gather() 的方法。任何想法？非常感谢！

另一个例子：

A = [[[1],     # A is (3, 5, 1)
       [2],
       [3],
       [4],
       [5]]],
     [[[10],
       [20],
       [30],
       [40],
       [50]]],
     [[[100],
       [200],
       [300],
       [400],
       [500]]]

B = [[1,4,3],     # B is (2,3)
     [2,3,5]]

C = [[[1, 2],     # C is (3, 3, 2)
       [4, 3],
       [3, 5]]],
     [[[10, 20],
       [40, 30],
       [30, 50]]],
     [[[100, 200],
       [400, 300],
       [300, 500]]]

Answer 1

B 张量的形状看起来不对，你的问题很难解析。但无论如何，TF 在这个问题上并不是很优雅。它需要一个非常特定的 B 形状。尝试类似于

import tensorflow as tf
import numpy as np

A = np.random.randn(10, 10, 7, 1).astype(np.float32)
A[0, 0, 1, 0] = 100001
A[0, 0, 3, 0] = 100002
A[0, 0, 5, 0] = 100003
A[0, 0, 2, 0] = 100004
A[0, 0, 4, 0] = 100005
A[0, 0, 6, 0] = 100006
A = tf.convert_to_tensor(A)

sess = tf.InteractiveSession()

B = np.array([
    [1, 3, 5],
    [2, 4, 6]
])

B = tf.convert_to_tensor(B)
B = tf.reshape(B, [-1])
B = tf.concat([tf.zeros_like(B), tf.zeros_like(B), B, tf.zeros_like(B)], axis=-1)
B = tf.reshape(B, [4, -1])
B = tf.transpose(B, [1, 0])
B = tf.reshape(B, [1, 2, 3, -1])


C = tf.gather_nd(A, B)
C = sess.run(C)
print C.shape
print C

输出是

[[[100001. 100002. 100003.]
  [100004. 100005. 100006.]]]