python - 如何使用另一个数组的元素作为索引对 tensorflow 中的张量进行切片？

Question

我正在寻找与 tf.unsorted_segment_sum 类似的函数，但我不想对段求和，我想将每个段作为张量。

例如，我有这段代码:(实际上，我有一个形状为 (10000, 63) 的张量，段数为 2500)

    to_be_sliced = tf.constant([[0.1, 0.2, 0.3, 0.4, 0.5],
                            [0.3, 0.2, 0.2, 0.6, 0.3],
                            [0.9, 0.8, 0.7, 0.6, 0.5],
                            [2.0, 2.0, 2.0, 2.0, 2.0]])

indices = tf.constant([0, 2, 0, 1])
num_segments = 3
tf.unsorted_segment_sum(to_be_sliced, indices, num_segments)

输出会在这里

array([sum(row1+row3), row4, row2]

我要找的是3个不同形状的张量(可能是张量列表)，第一个包含原始的第一行和第三行(形状为(2, 5))，第二个包含第4行(形状of (1, 5)), 第三个包含第二行，像这样:

[array([[0.1, 0.2, 0.3, 0.4, 0.5],
        [0.9, 0.8, 0.7, 0.6, 0.5]]),
 array([[2.0, 2.0, 2.0, 2.0, 2.0]]),
 array([[0.3, 0.2, 0.2, 0.6, 0.3]])]

提前致谢!

Answer 1

你可以这样做:

import tensorflow as tf

to_be_sliced = tf.constant([[0.1, 0.2, 0.3, 0.4, 0.5],
                            [0.3, 0.2, 0.2, 0.6, 0.3],
                            [0.9, 0.8, 0.7, 0.6, 0.5],
                            [2.0, 2.0, 2.0, 2.0, 2.0]])
indices = tf.constant([0, 2, 0, 1])
num_segments = 3
result = [tf.boolean_mask(to_be_sliced, tf.equal(indices, i)) for i in range(num_segments)]
with tf.Session() as sess:
    print(*sess.run(result), sep='\n')

输出:

[[0.1 0.2 0.3 0.4 0.5]
 [0.9 0.8 0.7 0.6 0.5]]
[[2. 2. 2. 2. 2.]]
[[0.3 0.2 0.2 0.6 0.3]]