我有一个类型为float的3D张量x
和一个类型为int的1D张量y
。我想得到x
的第二轴的每个切片的平均值,从0到对应于y
的每个元素的索引。换句话说,如果x
和y
是numpy数组,我希望
In [1]: y = [1, 2, 1, 1]
In [2]: x = np.array([[[1,2],[3,4],[5,6]], [[1,2],[3,4],[5,6]], [[1,2],[3,4],[5,6]], [[1,2],[3,4],[5,6]]])
In [3]: np.array([np.mean(x[index, :item], axis=0) for index, item in enumerate(y)])
Out[22]:
array([[ 1., 2.],
[ 2., 3.],
[ 1., 2.],
[ 1., 2.]])
最简单的方法是什么?
答案 0 :(得分:2)
在一般情况下,您可以使用tf.while_loop
:
import numpy as np
import tensorflow as tf
y = tf.constant([1, 2, 1, 1])
x = tf.constant([[[1,2],[3,4],[5,6]],
[[1,2],[3,4],[5,6]],
[[1,2],[3,4],[5,6]],
[[1,2],[3,4],[5,6]]], dtype=np.float32)
y_len = tf.shape(y)[0]
idx = tf.zeros((), dtype=tf.int32)
res = tf.zeros((0,2))
_, res = tf.while_loop(
lambda idx, res: idx < y_len,
lambda idx, res: (idx + 1, tf.concat([res, tf.reduce_mean(x[idx, :y[idx]], axis=0)[tf.newaxis]], axis=0)),
[idx, res],
shape_invariants=[idx.get_shape(), tf.TensorShape([None, 2])])
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
res.eval()
# returns
# array([[ 1., 2.],
# [ 2., 3.],
# [ 1., 2.],
# [ 1., 2.]], dtype=float32)
在不太常见的情况下,y
的长度在图形构造时已知,你可以在python中自己使用tf.while_loop
和循环(如果{可能会产生一个大图形) {1}}有很多元素。)
y
请注意,您也可以简单地级联更新,与y_len = y.shape.num_elements()
res = tf.Variable(np.zeros((y_len, 2), dtype=np.float32))
res = tf.tuple([res] + [res[idx].assign(tf.reduce_mean(x[idx, :y[idx]], axis=0))
for idx in range(y_len)])[0]
的一般情况不同:
tf.while_loop
但现在更新需要按顺序进行。在前一个解决方案中,每行发生的更新是独立的,可以并行运行,我认为这样更好。