Tensorflow 用 1d 张量索引到 2d 张量答案

【问题标题】：Tensorflow indexing into 2d tensor with 1d tensorTensorflow 用 1d 张量索引到 2d 张量
【发布时间】：2016-07-20 23:48:36
【问题描述】：

我有一个形状为 [batch_size, D] 的二维张量 A 和一个形状为 [batch_size] 的一维张量 B。 B 的每个元素都是A 的列索引，对于A 的每一行，例如。 B[i] in [0,D).

张量流中获取值A[B]的最佳方法是什么

例如：

A = tf.constant([[0,1,2],
                 [3,4,5]])
B = tf.constant([2,1])

具有所需的输出：

some_slice_func(A, B) -> [2,4]

还有另一个约束。在实践中，batch_size 实际上是None。

提前致谢！

【问题讨论】：

标签： tensorflow

【解决方案1】：

我能够使用线性索引使其工作：

def vector_slice(A, B):
    """ Returns values of rows i of A at column B[i]

    where A is a 2D Tensor with shape [None, D] 
    and B is a 1D Tensor with shape [None] 
    with type int32 elements in [0,D)

    Example:
      A =[[1,2], B = [0,1], vector_slice(A,B) -> [1,4]
          [3,4]]
    """
    linear_index = (tf.shape(A)[1]
                   * tf.range(0,tf.shape(A)[0]))
    linear_A = tf.reshape(A, [-1])
    return tf.gather(linear_A, B + linear_index)

这感觉有点hacky。

如果有人知道更好（如更清晰或更快），也请留下答案！（我暂时不会接受自己的）

【讨论】：

【解决方案2】：

@Eugene Brevdo 所说的代码：

def vector_slice(A, B):
    """ Returns values of rows i of A at column B[i]

    where A is a 2D Tensor with shape [None, D]
    and B is a 1D Tensor with shape [None]
    with type int32 elements in [0,D)

    Example:
      A =[[1,2], B = [0,1], vector_slice(A,B) -> [1,4]
          [3,4]]
    """
    B = tf.expand_dims(B, 1)
    range = tf.expand_dims(tf.range(tf.shape(B)[0]), 1)
    ind = tf.concat([range, B], 1)
    return tf.gather_nd(A, ind)

【讨论】：

【解决方案3】：

最简单的方法可能是通过连接 range(batch_size) 和 B 来构建适当的 2d 索引，以获得 batch_size x 2 矩阵。然后将其传递给 tf.gather_nd。

【讨论】：

【解决方案4】：

最简单的方法是这样做：

def tensor_slice(target_tensor, index_tensor):
    indices = tf.stack([tf.range(tf.shape(index_tensor)[0]), index_tensor], 1)
    return tf.gather_nd(target_tensor, indices)

【讨论】：

【解决方案5】：

考虑使用tf.one_hot、tf.math.multiply和tf.reduce_sum来解决。

例如

def vector_slice (inputs, inds, axis = None):
    axis = axis if axis is not None else tf.rank(inds) - 1

    inds = tf.one_hot(inds, inputs.shape[axis])
    for i in tf.range(tf.rank(inputs) - tf.rank(inds)):
        inds = tf.expand_dims(inds, axis = -1)
    
    inds = tf.cast(inds, dtype = inputs.dtype)
    
    x = tf.multiply(inputs, inds)

    return tf.reduce_sum(x, axis = axis)

【讨论】：