python - 使用整数张量在 tensorflow 中索引张量

Question

我的问题类似于 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].

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

Answer 1

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

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

基本上就是这样

1. 将索引数组展平为 1d，tf.reshape(p, (-1,));
2. 收集 mu[:,0] 中的元素(mu 的第一列)；
3. 然后将其 reshape 为 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)