python之如何从 tensorflow 张量中获取不同维度的切片。

我有一个浮点类型的 3D 张量 x 和一个 int 类型的一维张量 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.]])  

最简单的方法是什么？

请您参考如下方法:

在一般情况下，您可以使用 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 的长度在图形构造时已知，您可以省去使用 tf.while_loop 并在 python 中循环(如果 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 的一般情况不同:

y_len = y.shape.num_elements() 
res = tf.zeros((0,2)) 
for idx in range(y_len): 
  res = tf.concat([res, tf.reduce_mean(x[idx, :y[idx]], axis=0)[tf.newaxis]], axis=0) 

但现在更新需要按顺序进行。在前一种解决方案中，每一行的更新都是独立发生的，可以并行运行，我认为这样更好。