python - 如何使用 Tensorflow 获取张量中最大值的索引？

Question

我正在寻找使用 TensorFlow 1.14 获得张量最大值(等级 2 张量)的最佳和优化方法(无循环)的索引，在 Google Colab 中，我使用 TF2.0 尝试了此代码:

import tensorflow as tf
A = tf.constant([[0.2,0.8],[0.3,0.9],[0.4,0.7],[0.5,0.4]])
b = tf.math.argmax(A,0)
bb = b.numpy()

这里的最大值的索引是 [1,1]，但问题是我必须将轴作为输入，即使我更改了轴，它也不会给我正确的值。

Answer 1

不知道这是否是最好的方法，但我发现这个(来自 numpy documentation )可能对你有帮助:

import numpy as np

c = A.numpy()
np.unravel_index(np.argmax(c, axis=None), c.shape) # outputs (1,1)