如何找到张量对象中每一行的最大索引？答案

【问题标题】：How to find the max index for each row in a tensor object?如何找到张量对象中每一行的最大索引？
【发布时间】：2019-07-07 23:19:38
【问题描述】：

所以我正在创建一个 pytorch 模型，对于前向传递，我正在应用我的前向传递方法来获取包含每个类的预测分数的分数张量。这个张量的形状是 [100, 10]。现在，我想通过将其与包含实际分数的 y 进行比较来获得准确性。这个张量的形状是 [100]。为了比较两者，我将使用torch.mean(scores == y)，我会计算有多少是相同的。

问题是我需要转换分数张量，以便每一行只包含每行中最高值的索引。例如，如果张量看起来像这样，

tensor(
    [[0.3232, -0.2321, 0.2332, -0.1231, 0.2435, 0.6728],

    [0.2323, -0.1231, -0.5321, -0.1452, 0.5435, 0.1722],

    [0.9823, -0.1321, -0.6433, 0.1231, 0.023, 0.0711]]
)

然后我希望它被转换，使其看起来像这样。

tensor([5, 4, 0])

我怎么能这样做？

【问题讨论】：

标签： python python-3.x tensorflow pytorch

【解决方案1】：

将argmax 与所需的dim（又名轴）一起使用

a = tensor(
    [[0.3232, -0.2321, 0.2332, -0.1231, 0.2435, 0.6728],
    [0.2323, -0.1231, -0.5321, -0.1452, 0.5435, 0.1722],
    [0.9823, -0.1321, -0.6433, 0.1231, 0.023, 0.0711]]
)

a.argmax(1)
# tensor([ 5,  4,  0])

【讨论】：