如何在TensorFlow中沿維度選擇張量的元素（而非切片）

Question

我有一個3x2x4張量：

x = tf.reshape(tf.range(24), (3,2,4))

<tf.Tensor: id=1928, shape=(3, 2, 4), dtype=int64, numpy=
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])>

並且我想通過沿第3維索3x2將其減少到3x2 。 這是索引向量的樣子：

y = tf.constant(np.array([[0, 1, 0, 0], [0, 0, 1, 0], [1, 0, 0, 0]]))

<tf.Tensor: id=2093, shape=(3, 4), dtype=int64, numpy=
array([[0, 1, 0, 0],
       [0, 0, 1, 0],
       [1, 0, 0, 0]])>

所需的輸出是：

<tf.Tensor: id=2103, shape=(3, 2), dtype=int64, numpy=
array([[ 1,  5],
       [10, 14],
       [16, 20]])>

我嘗試了tf.batch_gather(x, y)但它給出了不同的輸出。 我需要collect_nd還是可以通過batch_gather解決？

Answer 1

您需要tf.boolean_mask() 。

import tensorflow as tf
import numpy as np

x = tf.reshape(tf.range(24), (3,2,4))
y = tf.constant(np.array([[0, 1, 0, 0], [0, 0, 1, 0], [1, 0, 0, 0]]))

result = tf.boolean_mask(tf.transpose(x,[0,2,1]),y)

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

[[ 1  5]
 [10 14]
 [16 20]]