跳到主要内容

numpy notes

np.argmax(logits, axis=-1) 是 NumPy 库中的一个函数，用于返回数组中指定轴上最大值的索引。下面是对这个函数的详细解释：

函数作用​

  • 功能：np.argmax 用于查找数组中最大元素的索引。
  • 参数：
    • logits：输入数组，可以是任意维度的 NumPy 数组。
    • axis：指定要查找最大值的轴。axis=-1 表示沿着最后一个轴（即最后一维）进行查找。

使用场景​

在深度学习和机器学习中，logits 通常指的是模型输出的未归一化的预测分数（例如，神经网络的输出层），这些分数通常用于分类任务。使用 np.argmax 可以帮助我们找到每个样本的预测类别。

示例​

假设我们有一个二维数组 logits，其中每一行代表一个样本的预测分数，每一列代表一个类别的分数：

import numpy as np

logits = np.array([[1.0, 2.0, 0.5],
                   [0.2, 0.8, 0.9],
                   [3.0, 1.0, 2.0]])

# 使用 np.argmax 查找每行的最大值索引
predicted_classes = np.argmax(logits, axis=-1)

print(predicted_classes)

输出​

运行上面的代码，predicted_classes 的值将是：

[1 2 0]

解释结果​

  • 第一行 [1.0, 2.0, 0.5] 的最大值是 2.0，索引为 1。
  • 第二行 [0.2, 0.8, 0.9] 的最大值是 0.9，索引为 2。
  • 第三行 [3.0, 1.0, 2.0] 的最大值是 3.0，索引为 0。

总结​

np.argmax(logits, axis=-1) 是一个非常有用的函数，特别是在分类任务中，它可以快速确定每个样本的预测类别。通过这种方式，我们可以将模型的输出转换为可解释的类别标签。