时间戳

首页 > 解决方案 > tensorflow 多维索引

python - tensorflow 多维索引

问题描述

我有

我想要做的很简单，用 3 个带有 numpy 的 for 循环：

y = np.random.randint(100,size=(5,10,20,3))
y_index= np.random.randint(3,size=(5,10,20))
y_slice = np.zeros_like(y_index)
for i in range(y.shape[0]):
    for j in range(y.shape[1]):
        for k in range(y.shape[2]):
            y_slice[i,j,k] = y[i,j,k,y_index[i,j,k]]
y_slice

我怎样才能在张量流中有效地做到这一点？我猜我需要使用 tf.gether_nd ...

标签： pythonnumpytensorflow

解决方案

您可以执行以下操作。基本上，首先您将除最后一个之外的所有维度展平，y并为 flatten 创建一个索引y。您进行索引，然后重塑为正确的形状。

y = tf.constant(np.random.normal(size=(5,10,20,3)), dtype='float32')
y_index = tf.constant(np.random.randint(3, size=(5,10,20)), dtype='int32')
# Creating an index like [(0,y_index[0]), (1, y_index[1]), ...]
inds = tf.stack([tf.range(5*10*20),tf.reshape(y_index,[-1])],axis=1)

y_slice = tf.reshape(tf.gather_nd(tf.reshape(y,[-1,3]),inds),[5,10,20])

推荐阅读