Python tensorflow中单个轴上指定切片的局部缩减
我正在尝试在二维阵列的单个轴上使用指定的切片执行局部缩减 我使用numpy的Python tensorflow中单个轴上指定切片的局部缩减,python,numpy,tensorflow,Python,Numpy,Tensorflow,我正在尝试在二维阵列的单个轴上使用指定的切片执行局部缩减 我使用numpy的numpy.ufunc.reduceat或numpy.add.reduceat实现了这一点,但我希望在tensorflow中也这样做,因为这个reduce操作的输入是tensorflow卷积的输出 我遇到了tf.math.reduce_sum,但我不确定如何在我的案例中使用它 如果我可以在tensorflow中执行reduceat操作,就像我可以利用GPU一样,那就太好了。您可以使用以下方法执行几乎相同的操作: 您可以使
numpy.ufunc.reduceat
或numpy.add.reduceat
实现了这一点,但我希望在tensorflow中也这样做,因为这个reduce操作的输入是tensorflow卷积的输出
我遇到了tf.math.reduce_sum
,但我不确定如何在我的案例中使用它
如果我可以在tensorflow中执行
reduceat
操作,就像我可以利用GPU一样,那就太好了。您可以使用以下方法执行几乎相同的操作:
您可以使用您想要使用的缩减功能替换上述内容。这与实际的唯一区别是特殊情况,索引[i]>=索引[i+1]
。posted函数要求对索引进行排序,如果出现索引[i]==索引[i+1]
的情况,则输出中相应的i
位置将为零,而不是a[index[i]
import tensorflow as tf
import numpy as np
def add_reduceat_tf(a, indices, axis=0):
a = tf.convert_to_tensor(a)
indices = tf.convert_to_tensor(indices)
# Transpose if necessary
transpose = not (isinstance(axis, int) and axis == 0)
if transpose:
axis = tf.convert_to_tensor(axis)
ndims = tf.cast(tf.rank(a), axis.dtype)
a = tf.transpose(a, tf.concat([[axis], tf.range(axis),
tf.range(axis + 1, ndims)], axis=0))
# Make segment ids
r = tf.range(tf.shape(a, out_type=indices.dtype)[0])
segments = tf.searchsorted(indices, r, side='right')
# Compute segmented sum and discard first unused segment
out = tf.math.segment_sum(a, segments)[1:]
# Transpose back if necessary
if transpose:
out = tf.transpose(out, tf.concat([tf.range(1, axis + 1), [0],
tf.range(axis + 1, ndims)], axis=0))
return out
# Test
np.random.seed(0)
a = np.random.rand(5, 10).astype(np.float32)
indices = [2, 4, 7]
axis = 1
# NumPy computation
out_np = np.add.reduceat(a, indices, axis=axis)
# TF computation
with tf.Graph().as_default(), tf.Session() as sess:
out = add_reduceat_tf(a, indices, axis=axis)
out_tf = sess.run(out)
# Check result
print(np.allclose(out_np, out_tf))
# True