使用整数张量在张量流中索引张量答案

【问题标题】：Indexing tensor in tensorflow using a tensor of integers使用整数张量在张量流中索引张量
【发布时间】：2018-06-09 11:47:15
【问题描述】：

我的问题与here 类似，但不完全相同。我有两个张量

mu: (shape=(1000,1), dtype=np.float32)
p : (shape=(100,30), dtype=np.int64)

我想要的是创建一个新的张量

x : (shape=(100,30), dtype=np.float32)

这样

x[i,j] = mu[p[i,j]]

这可以使用高级索引在 numpy 中完成

x = mu[p]

我曾尝试使用tf.gather_nd(mu, p) 命令，但在我的情况下，我收到以下错误

*** ValueError: indices.shape[-1] must be <= params.rank, but saw indices shape: [100,30] and params shape: [1000] for 'GatherNd_2' (op: 'GatherNd') with input shapes: [1000], [100,30].

因此，为了使用它，我必须建立一个新的坐标张量。有没有更简单的方法来完成我想要的？

【问题讨论】：

标签： python numpy tensorflow

【解决方案1】：

这是一个可行的解决方案：

tf.reshape(tf.gather(mu[:,0], tf.reshape(p, (-1,))), p.shape)

基本上是这样

将索引数组展平为 1d，tf.reshape(p, (-1,))；
收集来自mu[:,0]的元素（mu的第一列）；
然后将其重塑为p 的形状。

小例子：

import tensorflow as tf
tf.InteractiveSession()

mu = tf.reshape(tf.multiply(tf.cast(tf.range(10), tf.float32), 0.1), (10, 1))
mu.eval()
#array([[ 0.        ],
#       [ 0.1       ],
#       [ 0.2       ],
#       [ 0.30000001],
#       [ 0.40000001],
#       [ 0.5       ],
#       [ 0.60000002],
#       [ 0.69999999],
#       [ 0.80000001],
#       [ 0.90000004]], dtype=float32)

p = tf.constant([[1,3],[2,4],[3,1]], dtype=tf.int64)

tf.reshape(tf.gather(mu[:,0], tf.reshape(p, (-1,))), p.shape).eval()

#array([[ 0.1       ,  0.30000001],
#       [ 0.2       ,  0.40000001],
#       [ 0.30000001,  0.1       ]], dtype=float32)

另外两个选项使用gather_nd 而不进行整形：

tf.gather_nd(mu[:,0], tf.expand_dims(p, axis=-1)).eval()

#array([[ 0.1       ,  0.30000001],
#       [ 0.2       ,  0.40000001],
#       [ 0.30000001,  0.1       ]], dtype=float32)

tf.gather_nd(mu, tf.stack((p, tf.zeros_like(p)), axis=-1)).eval()

#array([[ 0.1       ,  0.30000001],
#       [ 0.2       ,  0.40000001],
#       [ 0.30000001,  0.1       ]], dtype=float32)

【讨论】：

【解决方案2】：

你可以使用tf.map_fn:

 x= tf.map_fn(lambda u: tf.gather(tf.squeeze(mu),u),p,dtype=mu.dtype)

map_fn 充当循环，在p 的第一个维度上运行，并且对于每个这样的切片都应用tf.gather。

【讨论】：