Question

我正在处理Numpy数组，代码如下所示：

z[np.arange(n), y]

其中z是2d数组，y是1d数组。此外，z.shape [0] == y.shape [0] == n。

我该如何做与TensorFlow张量等效的事情？

Answer 1

您可以使用tf.gather_nd来获取所需的索引。

import numpy as np
import tensorflow as tf

# Numpy implementation
n = 3
z = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
y = np.array([0, 1, 1])

assert z.shape[0] == y.shape[0] == n

np_out = z[np.arange(n), y]

# TF implementation
tf.reset_default_graph()

range_t = tf.range(n) # Equiv to np.arange
x_y = tf.stack([range_t, y], axis=1) # Get (x,y) as a tuple
pick_by_index_from_z = tf.gather_nd(z, x_y) # Pick the right values from z

with tf.Session() as sess:
  tf_out = sess.run(pick_by_index_from_z)

# The np and tf values should be the same
assert (np_out == tf_out).all()

print('z:')
print(z)
print('\nnp_out:')
print(np_out)
print('\ntf_out:')
print(tf_out)

这给出了输出：

z:
[[1 2 3]
 [4 5 6]
 [7 8 9]]

np_out:
[1 5 8]

tf_out:
[1 5 8]

TensorFlow中的Numpy代码相当于什么？

1 个答案: