从张量张量流中获取最大值

Question

a = np.array([[0, 10], [1.1, 0]])
# a.shape -> (2, 2)
arg_max = tf.math.argmax(input = a, axis=1)
# arg_max.shape -> (2, )

如何使用arg_maxlike [[10], [1.1]]从a获取最大值

Answer 1

如果没有arg_max，您可以使用：

a = np.array([[0, 10], [1.1, 0]])
tf.reduce_max(a, axis=1).numpy()
# Output
# array([10. ,  1.1])

Answer 2

对于数组，按索引获取值可能是这样的：

[[i[j]] for i,j in zip(a,arg_max)] #[[10.0], [1.1]]
# OR
[i[j] for i,j in zip(a,arg_max)]   #[10.0, 1.1]

不使用arg_max：

tf.nn.top_k(a,k=1)[0].numpy()