EDIT:
虽然下面的解决方案可能涵盖一些其他奇怪的用途,但只需转置数据就可以更轻松地解决此问题。事实证明,尽管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))
Output:
[[ 0.30000001 0.5 1.5 ]
[ 0.40000001 0.5 1.5 ]]
原帖:
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])
Output:
array([[ 0.3, 0.5, 1.5 ],
[ 0.4, 0.5, 1.5 ]], dtype=float32)