【学习笔记】torch.max()[]详解

torch.max(input, dim) 函数

output = torch.max(input, dim)

输入

• input是softmax函数输出的一个tensor
• dim是max函数索引的维度0/1，0是每列的最大值，1是每行的最大值

输出

• 函数会返回两个tensor，第一个tensor是每行的最大值；第二个tensor是每行最大值的索引。

在多分类任务中我们并不需要知道各类别的预测概率，所以返回值的第一个tensor对分类任务没有帮助，而第二个tensor包含了预测最大概率的索引，所以在实际使用中我们仅获取第二个tensor即可。

import torch
a = torch.tensor([[1,5,62,54], [2,6,2,6], [2,65,2,6]])
print(a)

 输出：

tensor([[ 1,  5, 62, 54],
        [ 2,  6,  2,  6],
        [ 2, 65,  2,  6]])

索引每行最大值：

torch.max(a, 1)

 输出：

torch.return_types.max(
values=tensor([62,  6, 65]),
indices=tensor([2, 3, 1]))

 在计算准确率时第一个tensor values是不需要的，所以我们只需提取第二个tensor，并将tensor格式的数据转换成array格式。

torch.max(a, 1)[1].numpy()

输出：

array([2, 3, 1], dtype=int64)

这样，我们就可以与标签值进行比对，计算模型预测准确率。

*注：在有的地方我们会看到torch.max(a, 1).data.numpy()的写法，这是因为在早期的pytorch的版本中，variable变量和tenosr是不一样的数据格式，variable可以进行反向传播，tensor不可以，需要将variable转变成tensor再转变成numpy。现在的版本已经将variable和tenosr合并，所以只用torch.max(a,1).numpy()就可以了。

准确率的计算

pred_y = torch.max(predict, 1)[1].numpy()
label_y = torch.max(label, 1)[1].data.numpy()
accuracy = (pred_y == label_y).sum() / len(label_y)

predict - softmax函数输出
label - 样本标签，这里假设它是one-hot编码