Question

正如this issue中正在进行广义切片一样，实现2D张量（矩阵）的op收集列的最佳方法是什么？例如，对于张量t：

1 2 3 4
5 6 7 8

和指数[1,3]，我想得到：

2 4
6 8

相当于numpy t[:, [1,3]]。

Answer 1

有一个名为tf.nn.embedding_lookup(params, ind) 的函数可以检索params张量的行。

为了达到你想要的效果，我们可以首先转置要从中选择某些列的张量t。然后查找tf.transpose(t)行（t列）。选择之后，我们将结果转换回来。

import tensorflow as tf


t = tf.constant([[1, 2, 3], 
                 [4, 5, 6]])
ind = tf.constant([0, 2])

result = tf.transpose(tf.nn.embedding_lookup(tf.transpose(t), ind))

with tf.Session() as sess:
    print(sess.run(result))

Answer 2

到目前为止，我通过展平输入并使用gather：

创建了一种解决方法

def gather_cols(params, indices, name=None):
    """Gather columns of a 2D tensor.

    Args:
        params: A 2D tensor.
        indices: A 1D tensor. Must be one of the following types: ``int32``, ``int64``.
        name: A name for the operation (optional).

    Returns:
        A 2D Tensor. Has the same type as ``params``.
    """
    with tf.op_scope([params, indices], name, "gather_cols") as scope:
        # Check input
        params = tf.convert_to_tensor(params, name="params")
        indices = tf.convert_to_tensor(indices, name="indices")
        try:
            params.get_shape().assert_has_rank(2)
        except ValueError:
            raise ValueError('\'params\' must be 2D.')
        try:
            indices.get_shape().assert_has_rank(1)
        except ValueError:
            raise ValueError('\'indices\' must be 1D.')

        # Define op
        p_shape = tf.shape(params)
        p_flat = tf.reshape(params, [-1])
        i_flat = tf.reshape(tf.reshape(tf.range(0, p_shape[0]) * p_shape[1],
                                       [-1, 1]) + indices, [-1])
        return tf.reshape(tf.gather(p_flat, i_flat),
                          [p_shape[0], -1])

适用于：

params = tf.constant([[1, 2, 3],
                      [4, 5, 6]])
indices = [0, 2]
op = gather_cols(params, indices)

产生预期的输出：

[[1 3]
 [4 6]]

Answer 3

与此同时，gather方法具有一个axis参数。

import tensorflow as tf
params = tf.constant([[1,2,3],[4,5,6]])
indices = [0,2]
op = tf.gather(params, indices, axis=1)

产生输出

[[1 3]
 [4 6]]

如何在TensorFlow中选择2D张量的某些列？

3 个答案: