在 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]]
言论:
- 所有列的索引将显示在一组索引中,且仅显示在一组索引中。
- 这必须在 Tensorflow 中完成,以便梯度可以通过此操作流动。
您知道如何执行该操作吗?我怀疑我需要使用 tf.slice,可能tf.while_loop。
你可以
用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]]
如果您
不介意使用 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]])