根据多组索引对2D张量的列求和

时间:2018-03-20 13:36:41

标签: python tensorflow

在tensorflow中,我想根据多组索引对2D张量的列进行求和。

例如:

总结以下张量的列

[[1 2 3 4 5]
 [5 4 3 2 1]]

根据2组指数(第一组设为总和0 1 2,第二组设为总和列3 4)

[[0,1,2],[3,4]]

应该提供2列

[[6  9]
 [12 3]]

备注:

  1. 所有专栏'指数将出现在一组且只有一组指数中。
  2. 这必须在Tensorflow中完成,以便渐变可以流经此操作。
  3. 您是否知道如何执行该操作?我怀疑我需要使用tf.slice,可能还需要使用tf.while_loop。

2 个答案:

答案 0 :(得分:2)

您可以使用tf.segment_sum

执行此操作
import tensorflow as tf

nums = [[1, 2, 3, 4, 5],
        [5, 4, 3, 2, 1]]
column_idx = [[0, 1, 2], [3, 4]]

with tf.Session() as sess:
    # Data as TF tensor
    data = tf.constant(nums)
    # Make segment ids
    segments = tf.concat([tf.tile([i], [len(lst)]) for i, lst in enumerate(column_idx)], axis=0)
    # Select columns
    data_cols = tf.gather(tf.transpose(data), tf.concat(column_idx, axis=0))
    col_sum = tf.transpose(tf.segment_sum(data_cols, segments))
    print(sess.run(col_sum))

输出:

[[ 6  9]
 [12  3]]

答案 1 :(得分:0)

如果您不介意用NumPy解决这个问题,我知道在NumPy中解决这个问题的粗暴方法。

import numpy as np

mat = np.array([[1, 2, 3, 4, 5], [5, 4, 3, 2, 1]])

grid1 = np.ix_([0], [0, 1, 2])
item1 = np.sum(mat[grid1])
grid2 = np.ix_([1], [0, 1, 2])
item2 = np.sum(mat[grid2])
grid3 = np.ix_([0], [3, 4])
item3 = np.sum(mat[grid3])
grid4 = np.ix_([1], [3, 4])
item4 = np.sum(mat[grid4])

result = np.array([[item1, item3], [item2, item4]])