简介
- 返回张量沿指定维度最大值的索引tf.argmax(张量名, axis = 操作轴)
代码
import numpy as np
import tensorflow as tf
test = np.array([[1, 2, 3], [2, 3, 4], [5, 4, 3], [8, 7, 2]])
print("test:\n", test)
print("每一列的最大值的索引:", tf.argmax(test, axis=0))
print("每一行的最大值的索引", tf.argmax(test, axis=1))