Tensorflow 用 1d 张量索引到 2d 张量 [英] Tensorflow indexing into 2d tensor with 1d tensor

问题描述

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

I have a 2D tensor A with shape [batch_size, D] , and a 1D tensor B with shape [batch_size]. Each element of B is a column index of A, for each row of A, eg. B[i] in [0,D).

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

What is the best way in tensorflow to get the values 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.

There is another constraint. In practice, batch_size is actually None.

提前致谢！

推荐答案

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

I was able to get it working using a linear index:

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.

This feels slightly hacky though.

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

If anyone knows a better (as in clearer or faster) please also leave an answer! (I won't accept my own for a while)

这篇关于Tensorflow 用 1d 张量索引到 2d 张量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！