将 gather_nd 与张量流中的聚合相结合

Question

我想对张量执行以下操作：

我想计算每行的平均值，但每行的值数量可以不同，如图所示。
我尝试了以下代码，但在 gather_nd 之后卡住了，我现在应该如何按照我想要的方式聚合它？

import tensorflow as tf

data = tf.constant(
    [
        [4, 4, 2, 4, 1],
        [6, 7, 6, 2, 4],
        [1, 0, 7, 6, 9],
        [5, 4, 1, 7, 4],
        [1, 2, 0, 3, 1]
    ]
)

indices = tf.constant([
    [0, 0], [0, 3],
    [1, 2],
    [2, 1], [2, 4],
    [3, 4],
    [4, 0], [4, 1], [4, 4]
])

print(tf.gather_nd(data, indices))
# --> tf.Tensor([4 4 6 0 9 4 1 2 1], shape=(9,), dtype=int32)
# now how to aggregate per row?

P.S。 - 这应该推广到更高的维度（我使用 2D 矩阵的原因是它更容易传达）。这意味着 data 将是 3D 而不是 2D，并且 indices 中的每个元素将包含 3 个项目而不是 2 个。

谢谢！

Answer 1

好的，用蒙版解决了。

import tensorflow as tf

data = tf.constant(
    [
        [4, 4, 2, 4, 1],
        [6, 7, 6, 2, 4],
        [1, 0, 7, 6, 9],
        [5, 4, 1, 7, 4],
        [1, 2, 0, 3, 1]
    ]
)

indices_no_batch = [
    [0, 3],
    [2],
    [1, 4],
    [4],
    [0, 1, 4]
]
batch_idxs = tf.repeat(tf.range(5), [len(row_size) for row_size in     indices_no_batch])
flat_list = [item for sublist in indices_no_batch for item in sublist]

index_to_data = tf.stack([batch_idxs, flat_list], axis=1)
candidate_mask = tf.scatter_nd(index_to_data, tf.ones(tf.shape(index_to_data)    [0]), tf.shape(data))
print(tf.reduce_mean(tf.ragged.boolean_mask(data, mask=tf.cast(candidate_mask, 'bool')), axis=1))
# --> tf.Tensor([4.         6.         4.5        4.         1.33333333], shape=(5,), dtype=float64)

归功于此答案 - how to average a tensor axis with specified mask in tensorflow

将 gather_nd 与张量流中的聚合相结合

combine gather_nd with aggregation in tensorflow

python

tensorflow

将 gather_nd 与张​​量流中的聚合相结合

combine gather_nd with aggregation in tensorflow

python

tensorflow

将 gather_nd 与张量流中的聚合相结合