【Python】tensorflow中的argmax()函数

在TensorFlow中，argmax() 函数是一个非常重要的操作，它用于返回给定张量（Tensor）沿指定轴的最大值的索引。这个函数在机器学习和深度学习应用中非常常见，尤其是在分类问题中，当我们需要确定哪个类别的预测概率最高时。

argmax() 函数的基本用法

argmax() 函数的一般形式如下：

tf.argmax(
    input,
    axis=None,
    name=None,
    dimension=None,  # 已弃用，请使用 axis
    output_type=tf.int64
)

• input：一个张量，表示要从中找出最大值的张量。
• axis：一个整数，指定要沿其找到最大值的轴。如果未指定，则默认对整个张量进行展平并返回单个最大值的索引。
• name：操作的名称（可选）。
• dimension：已弃用的参数，之前用于指定轴，现在应使用 axis。
• output_type：返回索引的数据类型，默认为 tf.int64。

示例

假设我们有一个二维张量，表示不同类别在不同样本上的预测概率：

import tensorflow as tf

# 创建一个二维张量，形状为 [3, 2]
predictions = tf.constant([[0.1, 0.9], [0.8, 0.2], [0.3, 0.7]], dtype=tf.float32)

# 沿着最后一个轴（axis=1）找到最大值的索引
class_indices = tf.argmax(predictions, axis=1)

# 创建一个 TensorFlow 会话并运行（在 TensorFlow 1.x 中需要这样做，TensorFlow 2.x 中通常不需要）
# with tf.Session() as sess:
#     print(sess.run(class_indices))

# 在 TensorFlow 2.x 中，可以直接运行
print(class_indices.numpy())  # 使用 .numpy() 方法将 TensorFlow 张量转换为 NumPy 数组（在 Eager Execution 模式下）

输出将是：

[1 0 1]

这表示第一个样本最可能的类别是索引为 1 的类别，第二个样本是索引为 0 的类别，第三个样本是索引为 1 的类别。

注意事项

• 在 TensorFlow 2.x 中，默认启用了 Eager Execution，因此你可以直接运行张量操作而无需创建会话。
• argmax() 函数返回的是最大值的索引，而不是最大值本身。
• 如果你的张量包含多个最大值（尽管这在大多数情况下不太可能，除非有特定的对称性或重复值），argmax() 函数将返回第一个找到的最大值的索引。
• 在处理分类问题时，通常会将 argmax() 函数应用于模型的输出（即预测概率），以确定每个样本最可能的类别。