Tensorflow unsorted_segment_sum 维度

Question

我正在使用 TensorFlow 的 tf.unsorted_segment_sum 方法，当我作为数据给出的张量只有一行时，它工作正常。例如：

tf.unsorted_segment_sum(tf.constant([0.2, 0.1, 0.5, 0.7, 0.8]),
                        tf.constant([0, 0, 1, 2, 2]), 3)

给出正确的结果：

array([ 0.3,  0.5 , 1.5 ], dtype=float32)

问题是，如果我使用包含多行的张量，我怎样才能获得每行的结果？例如，如果我尝试使用两行的张量：

tf.unsorted_segment_sum(tf.constant([[0.2, 0.1, 0.5, 0.7, 0.8],
                                     [0.2, 0.2, 0.5, 0.7, 0.8]]),
                        tf.constant([[0, 0, 1, 2, 2],
                                     [0, 0, 1, 2, 2]]), 3)

我期望的结果是：

array([ [ 0.3,  0.5 , 1.5 ], [ 0.4, 0.5, 1.5 ] ], dtype=float32)

但我得到的是：

array([ 0.7,  1. ,  3. ], dtype=float32)

我想知道是否有人知道如何在不使用 for 循环的情况下获取每一行的结果？

提前致谢

Answer 1

编辑：

虽然下面的解决方案可能涵盖一些额外的奇怪用途，但只要转置数据就可以更轻松地解决这个问题。事实证明，即使 tf.unsorted_segment_sum 没有 axis 参数，它也只能沿着一个轴工作，只要它是第一个轴。所以你可以这样做：

import tensorflow as tf

with tf.Session() as sess:
    data = tf.constant([[0.2, 0.1, 0.5, 0.7, 0.8],
                        [0.2, 0.2, 0.5, 0.7, 0.8]])
    idx = tf.constant([0, 0, 1, 2, 2])
    result = tf.transpose(tf.unsorted_segment_sum(tf.transpose(data), idx, 3))
    print(sess.run(result))

输出：

[[ 0.30000001  0.5         1.5       ]
 [ 0.40000001  0.5         1.5       ]]

原版POST:

tf.unsorted_segment_sum 不支持单轴工作。最简单的解决方案是将操作应用于每一行，然后将它们连接回去：

data = tf.constant([[0.2, 0.1, 0.5, 0.7, 0.8],
                    [0.2, 0.2, 0.5, 0.7, 0.8]])
segment_ids = tf.constant([[0, 0, 1, 2, 2],
                           [0, 0, 1, 2, 2]])
num_segments = 3
rows = []
for data_i, ids_i in zip(data, segment_ids):
    rows.append(tf.unsorted_segment_sum(data_i, ids_i))
result = tf.stack(rows, axis=0)

但是，这有缺点：1) 它只适用于静态形状的张量（也就是说，你需要有固定的行数）和 2) 它可能效率不高。第一个可以使用 tf.while_loop 来规避，但是，它会很复杂，而且它需要你一个接一个地连接行，这是非常低效的。另外，您已经说过要避免循环。

更好的选择是为每一行使用不同的 ID。例如，您可以向 segment_id 中的每个值添加类似 num_segments * row_index 的内容，这样您就可以保证每一行都有自己的一组 ID：

num_rows = tf.shape(segment_ids)[0]
rows_idx = tf.range(num_rows)
segment_ids_per_row = segment_ids + num_segments * tf.expand_dims(rows_idx, axis=1)

然后您可以应用操作和重塑以获得您想要的张量：

seg_sums = tf.unsorted_segment_sum(data, segment_ids_per_row,
                                   num_segments * num_rows)
result = tf.reshape(seg_sums, [-1, num_segments])

输出：

array([[ 0.3, 0.5, 1.5 ],
       [ 0.4, 0.5, 1.5 ]], dtype=float32)

Tensorflow unsorted_segment_sum 维度

Tensorflow unsorted_segment_sum dimension

tensorflow

sum

segment