如何按给定段收集张量元素？答案

【问题标题】：How to collect tensor elements by a given segment?如何按给定段收集张量元素？
【发布时间】：2019-06-27 17:45:20
【问题描述】：

我正在尝试实现“segment_collect”（非常类似于segment_max，但收集到张量而不是取最大值）。

t = tf.constant(["a", "b", "c", "d"])
s = tf.constant([0, 1, 1, 0])
r = tf.segment_collect(t, s)  # r == [["a", "d"], ["b", "c"]]

一个简单的实现是逐行构建结果，在以下伪代码中：

r = []
for i in range(2):
    mask = tf.equal(s, i)
    values = tf.boolean_mask(t, mask)
    r.append(values)
# convert r into a tensor at last

但这不是很有效。

一个后续问题是：是否有一种通用的方法来对张量进行分组/聚合？除了 tensorflow 中的 segment_{min/max/mean/prod/sum} 之外，这将允许更多的操作，例如 segment_size、segment_median、segment_percentile。

【问题讨论】：

标签： python tensorflow

【解决方案1】：

您可能会发现 tf.gather 和 tf.nn.topk 很有帮助：

tf.gather(t, tf.nn.top_k(-s, k=tf.shape(s)[0]).indices)

这适用于 TF 1.x 和 TF 2.0。如果需要，重塑结果：

tf.reshape(tf.gather(t, tf.nn.top_k(-s, k=tf.shape(s)[0]).indices), shape=(-1, 2))

当然，reshape 假设元素被分成大小相等的组（在本例中为两组）。

tf.keras.backend.eval(tf.gather(t, tf.nn.top_k(-s, k=tf.shape(s)[0]).indices))                                                           
# array([b'a', b'd', b'b', b'c'], dtype=object)

【讨论】：

【解决方案2】：

如果有人在这里寻找答案，tensorflow 已经为您提供了帮助：

t = tf.constant(["a", "b", "c", "d"]) # input tensor
s = tf.constant([0, 1, 1, 0])         # segment ids 
n = tf.unique(s).y.shape[0]           # number of segments

r = tf.dynamic_partition(t, s, n)     # r == [["a", "d"], ["b", "c"]]

如果段包含不同数量的元素，这也适用。

【讨论】：