Question

我正在尝试从3级tf.float张量x和2级{{1}构建2级tf.float张量y }}张量，tf.int32为：

z

我知道我需要将x[i][j] = y[i,z[i][j],j]用作：

tf.gather_nd

其中

x = tf.gather_nd(y,indices)

但是，我在使用tensorflow函数将indices[i][j][:] = [i,z[i][j],j]扩展到更高级别以构造z时遇到麻烦。

我正在尝试以向量化形式维护这些操作。

简单地使用indices作为

是否更实用？

tf.stack

其中

indices = tf.stack([ii,z,jj],axis=-1)

和

ii[i,:] = i？

Answer 1

我认为这可以满足您的需求：

import tensorflow as tf
import numpy as np

# Inputs
y = tf.placeholder(tf.float32, [None, None, None])
z = tf.placeholder(tf.int32, [None, None])
# Make first and last indices
y_shape = tf.shape(y)
ii, jj = tf.meshgrid(tf.range(y_shape[0]), tf.range(y_shape[2]), indexing='ij')
# Make full ND index
idx = tf.stack([ii, z, jj], axis=-1)
# Gather result
x = tf.gather_nd(y, idx)
# Test
with tf.Session() as sess:
    # Numbers from 0 to 11 in a (3, 4) matrix
    a = np.arange(12).reshape((3, 4))
    # Make Y with replicas of the matrix multiplied by 1, 10 and 100
    y_val = np.stack([a, a * 10, a * 100], axis=1).astype(np.float32)
    # Z will be a (3, 4) matrix of values 0, 1, 2, 0, 1, 2, ...
    z_val = (a % 3).astype(np.int32)
    # X  should have numbers from 0 to 11 multiplied by 1, 10, 100, 1, 10, 100, ...
    x_val = sess.run(x,  feed_dict={y: y_val, z: z_val}) #, feed_dict={y: y_val, z: z_val})
    print(x_val)

输出：

[[   0.   10.  200.    3.]
 [  40.  500.    6.   70.]
 [ 800.    9.  100. 1100.]]

使用collect_nd构建矩阵

1 个答案: